The question about the memory usage between gather-scatter and SpMM on RGCNConv

[zirui-ray-liu]

Hi,

From the official documents (https://pytorch-geometric.readthedocs.io/en/latest/notes/sparse_tensor.html), It seems that gather-scatter should has large memory footprint than SpMM based method. However, on RGCNConv, I found that SpMM consumes more memory than gather-scatter. I am wondering if this is a bug or not?

The test code is attached: (to apply scatter-gather, please add --use_longtensor)

import argparse
import os.path as osp
import os
#from memory_profiler import profile
#os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import pdb
import time
import numpy as np
import torch
import torch.nn.functional as F
from torch_geometric.datasets import Entities
from torch_geometric.utils import k_hop_subgraph
from torch_geometric.nn import RGCNConv, FastRGCNConv
import torch_geometric.transforms as T

parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, choices=['rgcn', 'qrgcn', 'fastrgcn'])
parser.add_argument('--dataset', type=str,
                    choices=['AIFB', 'MUTAG', 'BGS', 'AM'], default='AIFB')
parser.add_argument('--debug_mem', type=bool, default=True)
parser.add_argument('--use_longtensor', action='store_true')
parser.add_argument('--gpu', type=int, default=2)
args = parser.parse_args()

MB = 1024**2
GB = 1024**3


def get_memory_usage(gpu, print_info=False):
    """Get accurate gpu memory usage by querying torch runtime"""
    allocated = torch.cuda.memory_allocated(gpu)
    reserved = torch.cuda.memory_reserved(gpu)
    if print_info:
        print("allocated: %.2f MB" % (allocated / 1024 / 1024), flush=True)
        print("reserved:  %.2f MB" % (reserved / 1024 / 1024), flush=True)
    return allocated


def compute_tensor_bytes(tensors):
    """Compute the bytes used by a list of tensors"""
    if not isinstance(tensors, (list, tuple)):
        tensors = [tensors]
    ret = 0
    for x in tensors:
        if x.dtype in [torch.int64, torch.long]:
            ret += np.prod(x.size()) * 8
        if x.dtype in [torch.float32, torch.int, torch.int32]:
            ret += np.prod(x.size()) * 4
        elif x.dtype in [torch.bfloat16, torch.float16, torch.int16]:
            ret += np.prod(x.size()) * 2
        elif x.dtype in [torch.int8]:
            ret += np.prod(x.size())
        else:
            print(x.dtype)
            raise ValueError()
    return ret


# initialize configuration

if args.model =='fastrgcn':
    RGCNConv = FastRGCNConv
if args.model =='rgcn':
    RGCNConv = RGCNConv

path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'Entities')
dataset = Entities(path, args.dataset)
data = dataset[0]

# BGS and AM graphs are too big to process them in a full-batch fashion.
# Since our model does only make use of a rather small receptive field, we
# filter the graph to only contain the nodes that are at most 2-hop neighbors
# away from any training/test node.
node_idx = torch.cat([data.train_idx, data.test_idx], dim=0)
node_idx, edge_index, mapping, edge_mask = k_hop_subgraph(
    node_idx, 2, data.edge_index, relabel_nodes=True)

data.num_nodes = node_idx.size(0)
data.edge_index = edge_index
data.edge_type = data.edge_type[edge_mask]
data.train_idx = mapping[:data.train_idx.size(0)]
data.test_idx = mapping[data.train_idx.size(0):]
if not args.use_longtensor:
    data = T.ToSparseTensor()(data)
    data.adj_t = data.adj_t.set_value(data.edge_type)


class Net(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = RGCNConv(data.num_nodes, 16, dataset.num_relations,
                              num_bases=30)
        self.relu = torch.nn.ReLU()
        self.conv2 = RGCNConv(16, dataset.num_classes, dataset.num_relations,
                              num_bases=30)

    def forward(self, edge_index, edge_type):
        print(f'before conv1 mem: {torch.cuda.memory_allocated()/MB} MB')
        x = self.conv1(None, edge_index, edge_type)
        print(f'conv1 memory usage: {torch.cuda.memory_allocated()/MB} MB')
        x = self.relu(x)
        print(f'relu memory usage: {torch.cuda.memory_allocated()/MB} MB')
        x = self.conv2(x, edge_index, edge_type)
        print(f'conv2 memory usage: {torch.cuda.memory_allocated()/MB} MB')
        return F.log_softmax(x, dim=1)


torch.cuda.set_device(args.gpu)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#device = torch.device('cpu') if args.dataset == 'AM' else device
print(f'selected device: {device}')
model = Net()
model = model.to(device)

for name, child in model.named_children():
    print(f'{name}:\n{child}')

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

def train(data):
    data = data.to(device)
    model.train()
    optimizer.zero_grad()
    out = model(data.edge_index, data.edge_type)
    loss = F.nll_loss(out[data.train_idx], data.train_y)
    loss.backward()
    optimizer.step()
    return loss.item()


def debug_train(data):
    print("=========== Model and optimizer only ===========")
    optimizer.zero_grad()
    usage = get_memory_usage(args.gpu,True)
    print("=========== Load data to gpu ===========")
    data = data.to(device)
    if not args.use_longtensor:
        data.adj_t.fill_cache_()
    init_mem = get_memory_usage(args.gpu,True)
    data_mem = init_mem / MB - usage / MB
    #model.reset_parameters()
    model.train()
    print("=========== Before Backward ===========")
    if args.use_longtensor:
        out = model(data.edge_index, data.edge_type)
    else:
        out = model(data.adj_t, None)
    loss = F.nll_loss(out[data.train_idx], data.train_y)
    before_backward = get_memory_usage(args.gpu,True)
    act_mem = get_memory_usage(args.gpu, False) - init_mem - compute_tensor_bytes([loss, out])
    res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (before_backward / MB,
                                                                       data_mem,
                                                                       act_mem / MB)
    print(f'max allocated mem (MB): {torch.cuda.max_memory_allocated(0) / MB}')
    print(res)
    print("=========== After Backward ===========")
    loss.backward()
    optimizer.step()
    del loss, out
    after_backward = get_memory_usage(args.gpu,True)
    total_mem = before_backward + (after_backward - init_mem)
    res = "Total Mem: %.2f MB\tData Mem: %.2f MB\tAct Mem: %.2f MB" % (total_mem / MB,
                                                                       data_mem,
                                                                       act_mem / MB)
    print(res)
    print(f'max allocated mem (MB): {torch.cuda.max_memory_allocated(0) / MB}')
    exit()


@torch.no_grad()
def test():
    model.eval()
    pred = model(data.edge_index, data.edge_type).argmax(dim=-1)
    train_acc = pred[data.train_idx].eq(data.train_y).to(torch.float).mean()
    test_acc = pred[data.test_idx].eq(data.test_y).to(torch.float).mean()
    return train_acc.item(), test_acc.item()


if args.debug_mem:
    num_epochs = 5
else:
    num_epochs = 100


for epoch in range(1, num_epochs + 1):
    if args.debug_mem:
        loss = debug_train(data)
    else:
        loss = train()

[rusty1s]

I can indeed replicate that (although allocated memory is nearly identical (~55MB)). One explanation might be that the SparseTensor is modified upon each relation and will therefore need to recompute its cache for each iteration (computing rowptr and colptr).

[YijianLiu]

Hello, I want ask you one question about RGCN.
AIFB is Heterograph, but for the class entities to get AIFB graph, we just get relation type id.
But I really need to get relation type name for AIFB to know which edge is inverse, really need your help!

[rusty1s]

See #5689. Please don't duplicate your questions all over the place :)