SparseTensor memory footprint

[jmpark0808]

Hi everyone,

I've recently started using PyTorch Geometric and wanted to check if I was using memory efficient aggregation correctly.

I have two reproducible scripts, one for utilizing edge_index and another for SparseTensor.

Using Edge Index

import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
import torch.nn as nn
from torch_geometric.nn.conv import GATv2Conv
import torch_geometric.transforms as T
use_amp = False


    
class DummyModelNormal(nn.Module):
    def __init__(self, nfeat, nhid, dropout, nheads, nlayer):
        super(DummyModelNormal, self).__init__()
        self.linear = nn.Linear(nfeat, nhid)
        self.convs = nn.ModuleList([GATv2Conv(in_channels=nhid, out_channels=nhid,
                                                                          heads=nheads, dropout=dropout, edge_dim=None,
                                                                            concat=False) for _ in range(nlayer)])
        
    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.linear(x)
        for conv in self.convs:
            x = conv(x, edge_index=edge_index, edge_attr=None)
            
        return x


Normal_model = DummyModelNormal(26,16, 0, 8, 6)


device = 0
optimizer_type = torch.optim.Adam
Normal_model.cpu()

optimizer = optimizer_type(Normal_model.parameters(), lr=.001)
a = torch.cuda.memory_allocated(device)
Normal_model.to(device)
b = torch.cuda.memory_allocated(device)
model_memory = b - a
print(f'Normal Model memory: {model_memory/1e9}')


torch.manual_seed(0)
adj = torch.empty(625, 625).uniform_(0, 1)
edge_index = torch.bernoulli(adj)
edge_index = edge_index.nonzero().t().contiguous()
print(edge_index.size())
dummy_features = torch.randn(625, 26)

data = Data(x=dummy_features, edge_index=edge_index)
data_list = [data]*4
loader = DataLoader(data_list, batch_size=4)
batch = next(iter(loader))

output = Normal_model(batch.to(device))
c = torch.cuda.memory_allocated(device)

if use_amp:
    amp_multiplier = .5
else:
    amp_multiplier = 1
forward_pass_memory = (c - b)*amp_multiplier
print(f'Normal Forward pass memory: {forward_pass_memory/1e9}')
gradient_memory = model_memory
if isinstance(optimizer, torch.optim.Adam):
    o = 2
elif isinstance(optimizer, torch.optim.RMSprop):
    o = 1
elif isinstance(optimizer, torch.optim.SGD):
    o = 0
elif isinstance(optimizer, torch.optim.Adagrad):
    o = 1
else:
    raise ValueError("Unsupported optimizer. Look up how many moments are" +
        "stored by your optimizer and add a case to the optimizer checker.")
gradient_moment_memory = o*gradient_memory
total_memory = model_memory + forward_pass_memory + gradient_memory + gradient_moment_memory 
print(f'Normal Total memory: {total_memory/1e9}')

Which returns

Normal Model memory: 0.000113152
torch.Size([2, 195251])
Normal Forward pass memory: 7.76621824
Normal Total memory: 7.766670848

Using SparseTensor

import torch
from torch_geometric.data import Data
from torch_geometric.loader import DataLoader
import torch.nn as nn
from torch_geometric.nn.conv import GATv2Conv
from torch_sparse import SparseTensor
import torch_geometric.transforms as T
use_amp = False


class DummyModelSparse(nn.Module):
    def __init__(self, nfeat, nhid, dropout, nheads, nlayer):
        super(DummyModelSparse, self).__init__()
        self.linear = nn.Linear(nfeat, nhid)
        self.convs = nn.ModuleList([GATv2Conv(in_channels=nhid, out_channels=nhid,
                                                                          heads=nheads, dropout=dropout, edge_dim=None,
                                                                            concat=False) for _ in range(nlayer)])
        
    def forward(self, data):
        x, edge_index = data.x, data.adj_t

        x = self.linear(x)
        for conv in self.convs:
            x = conv(x, edge_index=edge_index, edge_attr=None)
            
        return x
    

Sparse_model = DummyModelSparse(26,16, 0, 8, 6)

device = 0
optimizer_type = torch.optim.Adam
Sparse_model.cpu()

optimizer = optimizer_type(Sparse_model.parameters(), lr=.001)
a = torch.cuda.memory_allocated(device)
Sparse_model.to(device)
b = torch.cuda.memory_allocated(device)
model_memory = b - a
print(f'Sparse Model memory: {model_memory/1e9}')

torch.manual_seed(0)
adj = torch.empty(625, 625).uniform_(0, 1)
edge_index = torch.bernoulli(adj)
edge_index = edge_index.nonzero().t().contiguous()
print(edge_index.size())

dummy_features = torch.randn(625, 26)

data = Data(x=dummy_features, edge_index=edge_index)
data_list = [data]*4
loader = DataLoader(data_list, batch_size=4)
batch = next(iter(loader))

transform = T.Compose([T.ToSparseTensor()])
batch = transform(batch)


output = Sparse_model(batch.to(device))
c = torch.cuda.memory_allocated(device)

if use_amp:
    amp_multiplier = .5
else:
    amp_multiplier = 1
forward_pass_memory = (c - b)*amp_multiplier
print(f'Sparse Forward pass memory: {forward_pass_memory/1e9}')
gradient_memory = model_memory
if isinstance(optimizer, torch.optim.Adam):
    o = 2
elif isinstance(optimizer, torch.optim.RMSprop):
    o = 1
elif isinstance(optimizer, torch.optim.SGD):
    o = 0
elif isinstance(optimizer, torch.optim.Adagrad):
    o = 1
else:
    raise ValueError("Unsupported optimizer. Look up how many moments are" +
        "stored by your optimizer and add a case to the optimizer checker.")
gradient_moment_memory = o*gradient_memory
total_memory = model_memory + forward_pass_memory + gradient_memory + gradient_moment_memory 
print(f'Sparse Total memory: {total_memory/1e9}')

which returns

Sparse Model memory: 0.000113152
torch.Size([2, 195251])
Sparse Forward pass memory: 7.811799552
Sparse Total memory: 7.81225216

Is it because the adjacency matrix (i.e. edge_index) is dense? or am I doing something wrong with the implementation?
Any help is greatly appreciated!

[rusty1s]

Not all GNN operators can leverage sparse-matrix multiplication inside their computations (and that's where the memory gains will come from). For example, if you switch to GCNConv or SAGEConv, you should definitely see memory footprint improvements.

[jmpark0808]

Thank you for the reply.

Based on the pytorch geometric cheat-sheet, it lists GATv2Conv support SparseTensor, but is that just saying that it supports the format but not necessarily the memory efficiency?

I quickly tried changing the layer to SAGEConv, but it doesn't seem like it's showing the memory improvements on my machine.
I changed

self.convs = nn.ModuleList([GATv2Conv(in_channels=nhid, out_channels=nhid,
                                                                           heads=nheads, dropout=dropout, edge_dim=None,
                                                                             concat=False) for _ in range(nlayer)])

to

self.convs = nn.ModuleList([SAGEConv(in_channels=nhid, out_channels=nhid) for _ in range(nlayer)])

and the output for the edge_index method is

Normal Model memory: 1.792e-05
torch.Size([2, 195251])
Normal Forward pass memory: 0.023528448
Normal Total memory: 0.023600128

and for the SparseTensor is

Sparse Model memory: 1.792e-05
torch.Size([2, 195251])
Sparse Forward pass memory: 0.02989568
Sparse Total memory: 0.02996736

[rusty1s]

Confirmed that SparseTensor uses message_and_aggregate and edge_index would use the message path in SAGEConv, so you will definitely avoid the blow-up in memory through materializing edge-level messages by using sparse-matrix multiplication. I don't really know why this isn't visible through torch.cuda.memory_allocated, as I am not sure if this is the most reliable way to track memory consumption.