Training GCN but using too many GPU memory

[RodyCoco]

Hi, guys, recently I meet some problem with my code because using too many GPU memory.
My input graph is a complete graph with 50 node, and each node having 66 features.
Batch size is set to 800.
Following is my code

import torch
import torch.nn as nn
import torch_geometric
import numpy as np

seed = 0
device_id = 0
N = 50

class GCN(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()

        self.conv1 = torch_geometric.nn.GCNConv(input_dim, hidden_dim[0], aggr='mean')
        self.conv2 = torch_geometric.nn.GCNConv(hidden_dim[0], hidden_dim[1], aggr='mean')
        self.conv3 = torch_geometric.nn.GCNConv(hidden_dim[1], output_dim, aggr='mean')
        self.Sigmoid = nn.Sigmoid()
        self.PReLU1 = nn.PReLU()
        self.PReLU2 = nn.PReLU()
        
    def forward(self, x, edge_index, edge_attr):

        x = self.conv1(x, edge_index, edge_attr)
        x = self.PReLU1(x)
        x = self.conv2(x, edge_index, edge_attr)
        x = self.PReLU2(x)
        x = self.conv3(x, edge_index, edge_attr)
        x = (self.Sigmoid(x)-0.5)*2
        return x

np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

model = GCN(
    input_dim=66,
    hidden_dim=[100,100],
    output_dim=2,
    ).cuda(device_id).double()
print(sum(p.numel() for p in model.parameters()))

x = torch.zeros((800, N, 66)).cuda(device_id).double()
y = torch.full((800, N, 2), 1.0).cuda(device_id).double()
from itertools import permutations
edge_index = torch.tensor(list(permutations(range(N), 2))).long().cuda(device_id).permute(1, 0)
edge_attr = torch.zeros((N*(N-1), 1)).cuda(device_id).double()
optimizer = torch.optim.Adam(list(model.parameters()), lr=1e-4)

output = model(x, edge_index, edge_attr)
loss_function = torch.nn.L1Loss() # can be replaced with any loss function
loss = loss_function(output, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(torch.cuda.memory_summary(device_id))

Environment

• PyG version: 2.0.3
• PyTorch version: 1.10.2
• Python version: 3.6.9
• CUDA/cuDNN version: 10.2
• How you installed PyTorch and PyG (conda, pip, source): pip
• Any other relevant information (e.g., version of torch-scatter):
• torch_spline_conv: 1.2.1
• torch_scatter: 2.0.9
• torch_sparse: 0.6.12
• torch_cluster: 1.5.9

I find that with above code, I using GPU memory with about 4.6GB. This quantity greatly exceeds my expectations. I would like to ask if this is normal in this situation? Thanks for your reading.

[EdisonLeeeee]

Why are you using double rather than float? This would lead to more memory overheads.

[RodyCoco]

Hi, I have update my torch_geometric, and I find a strange BUG, following is my testing code:

import torch
import torch.nn as nn
import torch_geometric
import numpy as np
from torch_sparse import SparseTensor

seed = 0
device_id = 0
N = 50

class GCN(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()

        self.conv1 = torch_geometric.nn.GCNConv(input_dim, hidden_dim[0], aggr='mean', cached=True)
        self.conv2 = torch_geometric.nn.GCNConv(hidden_dim[0], hidden_dim[1], aggr='mean', cached=True)
        self.conv3 = torch_geometric.nn.GCNConv(hidden_dim[1], output_dim, aggr='mean', cached=True)
        self.Sigmoid = nn.Sigmoid()
        self.PReLU1 = nn.PReLU()
        self.PReLU2 = nn.PReLU()
        
    def forward(self, x, adj_t, edge_attr):

        x = self.conv1(x, adj_t, edge_attr)
        x = self.PReLU1(x)
        x = self.conv2(x, adj_t, edge_attr)
        x = self.PReLU2(x)
        x = self.conv3(x, adj_t, edge_attr)
        x = (self.Sigmoid(x)-0.5)*2
        return x

np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

model = GCN(
    input_dim=66,
    hidden_dim=[100,100],
    output_dim=2,
    ).cuda(device_id).double()
print(sum(p.numel() for p in model.parameters()))

x = torch.zeros((800, N, 66)).cuda(device_id).double()
y = torch.full((800, N, 2), 1.0).cuda(device_id).double()
from itertools import permutations
edge_index = torch.tensor(list(permutations(range(N), 2))).long().cuda(device_id).permute(1, 0)
edge_attr = torch.zeros((N*(N-1), 1)).cuda(device_id).double()
optimizer = torch.optim.Adam(list(model.parameters()), lr=1e-4)
adj = SparseTensor(row=edge_index[0], col=edge_index[1])
from torch_geometric.data import Data
graph = Data(x=x, adj_t=adj.t(), edge_index=edge_index, edge_attr=edge_attr).cuda(device_id)

for i in range(1, 11):
    output = model(graph.x, graph.adj_t, edge_attr)
    loss_function = torch.nn.L1Loss() # can be replaced with any loss function
    loss = loss_function(output, y)
    optimizer.zero_grad()
    loss.backward()
    print(f"epoch: {i}, loss: {loss}")
    optimizer.step()

If I change line 55 output = model(graph.x, graph.adj_t, edge_attr) into output = model(graph.x, graph.edge_index, edge_attr)
, the process of loss change is not the same.
I have test the example code in https://pytorch-geometric.readthedocs.io/en/latest/advanced/sparse_tensor.html

import torch
import torch.nn.functional as F

from torch_geometric.nn import GCNConv
import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
import numpy as np

seed = 1
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

dataset = Planetoid("Planetoid", name="Cora", transform=T.ToSparseTensor())
data = dataset[0]

class GNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16, cached=True)
        self.conv2 = GCNConv(16, dataset.num_classes, cached=True)

    def forward(self, x, adj_t):
        x = self.conv1(x, adj_t)
        x = F.relu(x)
        x = self.conv2(x, adj_t)
        return F.log_softmax(x, dim=1)

model = GNN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train(data):
    model.train()
    optimizer.zero_grad()
    row, col, edge_attr = data.adj_t.t().coo()
    edge_index = torch.stack([row, col], dim=0)
    out = model(data.x, data.adj_t)
    loss = F.nll_loss(out, data.y)
    loss.backward()
    optimizer.step()
    return float(loss)

for epoch in range(1, 10):
    loss = train(data)
    print(loss)

If I change line 38
out = model(data.x, data.adj_t)
into
out = model(data.x, edge_index)
, the process of loss change is the same.
It seems that if I add edge_attr into GCN input, the result will be different when using adj_t and edge_index as GCN input respectivly. Is it normal? I have problem when training GCN network for my project, the problem is when I using adj_t and edge_index (with edge_attr as input), their loss change is not the same, and the former converge, but the latter is not converge.
So I wonder that is I correctly using the SparseTensor or there is some bug? Thanks for your reply.

[rusty1s]

Yes, that's normal. SparseTensor expects that edge features are included in it. So you may need to write:

adj = SparseTensor(row=edge_index[0], col=edge_index[1], value=edge_attr)

conv(x, adj.t())

[RodyCoco]

I know GMMconv can write in the way you provided, but when I using GCNconv, here comes the error: RuntimeError: output with shape [2500] doesn't match the broadcast shape [2500, 2500]
The adj matrix received in GCNconv is expected to be 2 dimension, but after add "value=edge_attr" in SparseTensor, it becomes a 3D tensor. Is there a way to reduce 3D SparseTensor to 2D SparseTensor? My 3D SparseTensor has the size (50, 50, 1), I want to resize it into (50, 50). I have search the use of SparseTensor, but no relevant information was found.

[rusty1s]

You may want to squeeze the value of SparseTensor to be one-dimensional:

SparseTensor(..., value=edge_attr.view(-1))

[RodyCoco]

This works for me, thanks for your reply.