Thanks for the great work. I am trying to check the memory usage based on GAS, I used APPNP which is provided to run on some large scale datasets, the following is the code, I just added several lines for memory usage within the forward:
class APPNP(ScalableGNN):
def __init__(self, num_nodes: int, in_channels, hidden_channels: int,
out_channels: int, num_mlp:int, num_layers: int, alpha: float,
dropout: float = 0.0, pool_size: Optional[int] = None,
buffer_size: Optional[int] = None, device=None):
super().__init__(num_nodes, out_channels, num_layers, pool_size,
buffer_size, device)
self.in_channels = in_channels
self.out_channels = out_channels
self.alpha = alpha
self.dropout = dropout
self.lins = torch.nn.ModuleList()
self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
self.bns = torch.nn.ModuleList()
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
for _ in range(num_mlp):
self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
self.lins.append(torch.nn.Linear(hidden_channels, out_channels))
self.reg_modules = self.lins[:1]
self.nonreg_modules = self.lins[1:]
def reset_parameters(self):
super().reset_parameters()
for lin in self.lins:
lin.reset_parameters()
for bn in self.bns:
bn.reset_parameters()
def forward(self, x: Tensor, adj_t: SparseTensor, *args) -> Tensor:
print("allocated1: %.2f MB" % (torch.cuda.memory_allocated() / 1024 / 1024), flush=True)
x = self.lins[0](x).relu_()
x = F.dropout(x, p=self.dropout, training=self.training)
for i, lin in enumerate(self.lins[1:-1]):
x = lin(x)
x = self.bns[i](x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.lins[-1](x)
x_0 = x[:adj_t.size(0)]
print("allocated2: %.2f MB" % (torch.cuda.memory_allocated() / 1024 / 1024), flush=True)
for history in self.histories:
x = (1 - self.alpha) * (adj_t @ x) + self.alpha * x_0
x = self.push_and_pull(history, x, *args)
x = (1 - self.alpha) * (adj_t @ x) + self.alpha * x_0
print("allocated3: %.2f MB" % (torch.cuda.memory_allocated() / 1024 / 1024), flush=True)
return x
@torch.no_grad()
def forward_layer(self, layer, x, adj_t, state):
if layer == 0:
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.lins[0](x)
x = x.relu()
x = F.dropout(x, p=self.dropout, training=self.training)
x = x_0 = self.lins[1](x)
state['x_0'] = x_0[:adj_t.size(0)]
x = (1 - self.alpha) * (adj_t @ x) + self.alpha * state['x_0']
return x
def mini_train(model, loader, criterion, optimizer, max_steps, grad_norm=None,
edge_dropout=0.0):
model.train()
total_loss = total_examples = 0
for i, (batch, batch_size, *args) in enumerate(loader):
x = batch.x.to(model.device)
adj_t = batch.adj_t.to(model.device)
y = batch.y[:batch_size].to(model.device)
train_mask = batch.train_mask[:batch_size].to(model.device)
if train_mask.sum() == 0:
continue
# We make use of edge dropout on ogbn-products to avoid overfitting.
adj_t = dropout(adj_t, p=edge_dropout)
optimizer.zero_grad()
out = model(x, adj_t, batch_size, *args)
loss = criterion(out[train_mask], y[train_mask])
print("middle: %.2f MB" % (torch.cuda.memory_allocated() / 1024 / 1024), flush=True)
loss.backward(retain_graph=True)
if grad_norm is not None:
torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
optimizer.step()
print("after: %.2f MB" % (torch.cuda.memory_allocated() / 1024 / 1024), flush=True)
total_loss += float(loss) * int(train_mask.sum())
total_examples += int(train_mask.sum())
# We may abort after a fixed number of steps to refresh histories...
if (i + 1) >= max_steps and (i + 1) < len(loader):
break
#print("allocated: %.2f MB" % (torch.cuda.memory_allocated() / 1024 / 1024), flush=True)
return total_loss / total_examples
@torch.no_grad()
def full_test(model, data):
model.eval()
return model(data.x.to(model.device), data.adj_t.to(model.device)).cpu()
@torch.no_grad()
def mini_test(model, loader):
model.eval()
return model(loader=loader)
Loading data... Done! [4.13s]
Computing METIS partitioning with 150 parts... Done! [72.42s]
Permuting data... Done! [52.20s]
Adding self-loops... Done! [3.61s]
Normalizing data... Done! [1.83s]
Pre-processing subgraphs... Done! [5.54s]
Pre-processing subgraphs... Done! [23.81s]
Calculating buffer size... Done! [0.01s] -> 104781
Fill history... allocated1: 0.00 MB
allocated2: 0.00 MB
allocated3: 4259.25 MB
Done! [3.61s]
I am wondering why the "allocated2" which should be the memory after MLP is 0, and there is a hugh increase between "allocated2" and "allocated3". And during training, the memory keeps increased and then becomes stable, is it normal? Thanks for the help.