Reaching an error with nodeneighborhood loader

[akul-goyal]

Hi,

I am getting the following error:

RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)

If I change the batch size of the loader then the error goes away. Investigating a little deeper, some of the subgraphs getting passed into the model are a couple of magnitudes bigger than the norm (like 6 nodes vs 300 nodes with batch size 3). The error spits out with these bigger subgraphs. I don't get this issue with high batch sizes (2048, 4096) but it appears in lower batch sizes (3->10).

[rusty1s]

Thanks for reporting. Do you have a minimal script to reproduce this? I'm happy to look into this.

[akul-goyal]

Hey Rusty, I will try to get you a simple script soon!

[akul-goyal]

import torch
from tqdm import tqdm
import torch.nn.functional as F
import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
import os.path as osp
import copy
import statistics
from torch_geometric.nn import GCNConv, VGAE, GATConv
from torch.nn import Sequential as Seq, Linear, ReLU
from torch.nn import GRUCell, Linear
from torch import nn
import pickle
import torch_geometric
import wandb
import torch_geometric.loader as tgLoader

torch_geometric.seed.seed_everything(100)
gEpochs = 5
numHeads = 8

class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, file=None, proc=None):
        super(Net, self).__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels, cached=True)
        self.conv2 = GCNConv(hidden_channels, hidden_channels, cached=True)
        self.conv3 = GCNConv(hidden_channels, out_channels, cached=True)
        self.GAT = GATConv(hidden_channels, hidden_channels, add_self_loops=False, dropout=0.8, negative_slope=0.01, heads =numHeads)
        self.GAT2 = GATConv(hidden_channels, hidden_channels, add_self_loops=False, negative_slope=0.01, dropout=0.8, heads = numHeads)
        self.GAT3 = GATConv(hidden_channels, hidden_channels, add_self_loops=False, negative_slope=0.01,dropout=0.6,  heads=1)
        self.GRU = GRUCell(hidden_channels*numHeads, hidden_channels)
        self.GRU3= GRUCell(hidden_channels * numHeads, hidden_channels)
        self.GRU2 = GRUCell(hidden_channels, hidden_channels)
        self.lin1 = Linear(in_channels, hidden_channels)
        self.lin2 = Linear(hidden_channels, out_channels)
        self.mlp = Seq(Linear((numHeads+1) * hidden_channels, hidden_channels),
                       ReLU(),
                       Linear(hidden_channels, hidden_channels))
        self.mlp2 = Seq(Linear(2 * hidden_channels, hidden_channels),
                       ReLU(),
                       Linear(hidden_channels, hidden_channels))
        self.allLin = Linear(hidden_channels, out_channels)


    def forward(self, x, edge_index, root_index = None, type = None, getWeights=False):
        # featurize
        encodeX = None
        x = nn.ReLU()(self.lin1(x))
        # 1st hop
        h = F.leaky_relu(self.GAT(x, edge_index))
        x = nn.ReLU()(self.GRU(h, x))
        # 2nd hop
        h = F.leaky_relu(self.GAT(x, edge_index))
        x = nn.ReLU()(self.GRU(h, x))
        # pick between emebedding (test)
        # get class distribution (train)
        return x

def createLoader(mask, data, batch = 4096):
    loader = tgLoader.NeighborLoader(data, num_neighbors=[-1, -1],
                                                input_nodes=mask, batch_size=batch, shuffle=False)
    return loader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
path = osp.join(osp.dirname(osp.realpath(__file__)), '../data/Planetoid')
dataset = Planetoid(path, name='Cora')
data = dataset[0]
trainLoader  = createLoader(data.train_mask, data)
valLoader = createLoader(data.val_mask, data, batch_size=3)
use_gpu = True if torch.cuda.is_available() else False
model = Net(data.x.shape[1], hidden_channels=128, out_channels=len(classes)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.00005)
criterion = nn.CrossEntropyLoss()

def run_loader(loader, train=True, saveZ = False):
    if train:
        model.train()
    else:
        model.eval()
    total_loss = 0
    total_acc = 0
    count = 0
    if saveZ:
        globalZ = torch.FloatTensor([])
    iter_loader = iter(loader)
    z_data = next(iter_loader, None)
    while z_data is not None:
        z_data = z_data.to(device)
        z, nun = model(z_data.x, z_data.edge_index)
        if saveZ:
            globalZ = torch.cat((globalZ, z[:z_data.batch_size].detach().to('cpu')),0)
        y_max_scores, y_pred = z.max(dim=1)
        acc = (z_data.y == y_pred).sum() / z_data.y.size(0)
        loss = criterion(z, z_data.y)
        total_loss += loss.item()
        total_acc += acc
        count += 1
        if train:
            loss.backward()
            optimizer.step()
        z_data = next(iter_loader, None)
    if saveZ:
        return total_loss/count, total_acc/count, globalZ
    return total_loss/count, total_acc/count
def run():
    global parDir
    pbar = tqdm(total=gEpochs)
    patience = 1000
    best_loss_val = 100
    for epoch in range(0, gEpochs):
        pbar.set_description(f'Training epoch: {epoch:04d}')
        train_loss, train_acc = run_loader(trainLoader, train=True)
        val_loss, acc = run_loader(valLoader, train=False)
        print(f'Val loss: {val_loss}, Val Acc: {acc}')
        pbar.update(1)
    pbar.close()
    model.load_state_dict(torch.load('modelStateDict.pkl'))
    _, _, z = run_loader(testLoader, train=False, saveZ=True)
    return z
run()

[akul-goyal]

The problem occurs when the batch size is small

[rusty1s]

Thanks for the script! I tried to run it but it complains about createLoader not returning a tuple:

trainLoader, trainLoaderRAdj, _ = createLoader(data.train_mask, data)
ValueError: not enough values to unpack (expected 3, got 1)

[akul-goyal]

@rusty1s Hey this is probably super late, but I fixed my code above. Let me know if it runs and brings out the bug. I can update the code again if you have more issues!

[akul-goyal]

@rusty1s Do you know if there was there any update on this?

[rusty1s]

Can you clarify? Last time your code is fixed :)