Neighborhood Sampler for AGNNCONV, TAGCONV etc

[JindeShubhamA]

Hello,

I was benchmarking the result of some open-source datasets on different Graph Neural Networks. I wanted to compare the performance of GCNConv, SAGEConv, GATConv, AGNNConv, TAGConv on the citation, co-author, pattern, and cluster dataset for the node classification task. Successfully got the results on Citation dataset :). When I tried on the co-author dataset which is much bigger than a citation, I get the 'CUDA out of memory' issue. I was able to get the results by implementing a Neighborhood Sampler for SAGEConv and GATConv using ' examples/reddit.py', and 'ogbn_products_sage.py'. For AGNNConv I get the following error

/torch_geometric/nn/conv/agnn_conv.py", line 63, in forward
num_nodes=x.size(self.node_dim))
AttributeError: 'tuple' object has no attribute 'size'

Can anyone suggest on how should I run AGNNConvm TAGConv using Neighborhood Sampler or any equivalent mechanism. I am also attaching my code with this. Also any advice in general on how should I compare the performance of the GNNs mentioned above on pattern and cluster dataset which are much more bigger in size. Any help will be appreciated.

Below is the code

import os.path as osp

import torch
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np
from torch_geometric.data import NeighborSampler
from torch_geometric.nn import SAGEConv, AGNNConv, GCNConv, TAGConv
from torch_geometric.datasets import Planetoid, Coauthor
# root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'products')
dataset = Coauthor(root='/tmp/Coauthor', name='physics')
# split_idx = dataset.get_idx_split()
# evaluator = Evaluator(name='ogbn-products')
data = dataset[0]

#### Added to create train ans test mask on 80-20 split on Coauthor dataset.
train_mask = 27587
num_nodes = 34493

train_mask_lst = [True]*num_nodes
for i in range(train_mask,num_nodes):
    train_mask_lst[i] = False
train_mask_lst = np.array(train_mask_lst)
train_mask_lst = torch.from_numpy(train_mask_lst) 

train_loader = NeighborSampler(data.edge_index, node_idx=train_mask_lst,
                               sizes=[15, 10, 5], batch_size=1024,
                               shuffle=True, num_workers=12)
subgraph_loader = NeighborSampler(data.edge_index, node_idx=None, sizes=[-1],
                                  batch_size=4096, shuffle=False,
                                  num_workers=12)


class Net(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
        super(Net, self).__init__()

        self.num_layers = num_layers

        self.convs = torch.nn.ModuleList()
        self.convs.append(AGNNConv(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(AGNNConv(hidden_channels, hidden_channels))
        self.convs.append(AGNNConv(hidden_channels, out_channels))

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()

    def forward(self, x, adjs):
        # `train_loader` computes the k-hop neighborhood of a batch of nodes,
        # and returns, for each layer, a bipartite graph object, holding the
        # bipartite edges `edge_index`, the index `e_id` of the original edges,
        # and the size/shape `size` of the bipartite graph.
        # Target nodes are also included in the source nodes so that one can
        # easily apply skip-connections or add self-loops.
        for i, (edge_index, _, size) in enumerate(adjs):
            x_target = x[:size[1]]  # Target nodes are always placed first.
            x = self.convs[i]((x, x_target), edge_index)
            if i != self.num_layers - 1:
                x = F.relu(x)
                x = F.dropout(x, p=0.5, training=self.training)
        return x.log_softmax(dim=-1)

    def inference(self, x_all):
        pbar = tqdm(total=x_all.size(0) * self.num_layers)
        pbar.set_description('Evaluating')

        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        total_edges = 0
        for i in range(self.num_layers):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to(device)
                total_edges += edge_index.size(1)
                x = x_all[n_id].to(device)
                x_target = x[:size[1]]
                x = self.convs[i]((x, x_target), edge_index)
                if i != self.num_layers - 1:
                    x = F.relu(x)
                xs.append(x.cpu())

                pbar.update(batch_size)

            x_all = torch.cat(xs, dim=0)

        pbar.close()

        return x_all


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model =Net(dataset.num_features, 32, dataset.num_classes, num_layers=3)
model = model.to(device)

x = data.x.to(device)
y = data.y.squeeze().to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
def train(epoch):
    model.train()

    pbar = tqdm(total=train_mask)
    pbar.set_description(f'Epoch {epoch:02d}')

    total_loss = total_correct = 0
    for batch_size, n_id, adjs in train_loader:
        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
        adjs = [adj.to(device) for adj in adjs]

        optimizer.zero_grad()
        out = model(x[n_id], adjs)
        loss = F.nll_loss(out, y[n_id[:batch_size]])
        loss.backward()
        optimizer.step()

        total_loss += float(loss)
        total_correct += int(out.argmax(dim=-1).eq(y[n_id[:batch_size]]).sum())
        pbar.update(batch_size)

    pbar.close()

    loss = total_loss / len(train_loader)
    approx_acc = total_correct / train_mask

    return loss, approx_acc


@torch.no_grad()
def test():
    model.eval()

    out = model.inference(x)

    y_true = y.cpu().unsqueeze(-1)
    y_pred = out.argmax(dim=-1)

    # results = []
    # for mask in [data.train_mask, data.val_mask, data.test_mask]:
    #     results += [int(y_pred[mask].eq(y_true[mask]).sum()) / int(mask.sum())]

    correct = y_pred[train_mask:].eq(data.y[train_mask:])


    print("Correct", correct)
    print("Size of Correct", len(correct))
    print("Sum correct", correct.sum())
    correct = correct.sum()
    print(len(y_pred[train_mask:]))
    print('Correct ', correct)
    print('Nodes ', num_nodes-train_mask)
    acc = correct / int(num_nodes-train_mask)
    print('Accuracy: {:.4f}'.format(acc))
      
    return acc


test_accs = []
for epoch in range(1, 200):
    loss, acc = train(epoch)
    print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {acc:.4f}')
test_acc = test()


test_acc = torch.tensor(test_accs)

[rusty1s]

There are indeed a few GNN operators in PyG that do not work in conjunction with NeighborSampler, e.g., AGNNConv. Those can be usually identified by the fact that they do not support a tuple of input node features or the size argument in forward.

In general, one does not need to make use of NeighborSampling for those small datasets. If you run into memory-issues on those datasets, I suggest you to look into working with SparseTensors. You can find a tutorial here.

[JindeShubhamA]

I tried cross_entropy still facing the same problem

[rusty1s]

This is weird, cross_entropy should work just fine. It expects a tensor of [num_nodes, num_classes] for predictions and a tensor of [num_nodes] for ground-truth labels.

[JindeShubhamA]

It expects a tensor of [num_nodes, num_classes] for predictions and a tensor of [num_nodes] for ground-truth labels.
Yes, it is correct. When I printed y_prediction, it is always coming as [1 0] for all the cases. It is not predicting anything for second class although there are enough examples of both classes. Also, I checked the neuron's activation before applying softmax, and neurons are only activated for the first class and not the second which clearly means the network is not learning activations for the second class.

I will attach my code for any bug in the implementation and maybe could also be run to cross-check the result.

mport os.path as osp

import torch
# from torch.nn import cross_entropy
import torch.nn.functional as F
from torch_geometric.datasets import PPI, GNNBenchmarkDataset
from torch_geometric.data import DataLoader
from torch_geometric.nn import GATConv, GCNConv, SAGEConv
from sklearn.metrics import f1_score
import numpy as np 

path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'PPI')
train_dataset = GNNBenchmarkDataset(path, name='PATTERN',  split='train')
val_dataset = GNNBenchmarkDataset(path, name='PATTERN', split='val')
test_dataset = GNNBenchmarkDataset(path, name='PATTERN',split='test')
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)

print(f'Dataset: {train_dataset}:')
print('======================')
print(f'Number of graphs: {len(train_dataset)}')
print(f'Number of features: {train_dataset.num_features}')
print(f'Number of classes: {train_dataset.num_classes}')


class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = GCNConv(train_dataset.num_node_features, 256)
        self.conv2 = GCNConv(256, 256)
        self.conv3 = GCNConv(256, train_dataset.num_classes)


    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = self.conv2(x, edge_index)
        x = self.conv3(x, edge_index)
        x = F.softmax(x,dim=1)
        return x

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.005)

def train():
    model.train()

    total_loss = 0
    for data in train_loader:
        data = data.to(device)
        optimizer.zero_grad()
        y_pred  = model(data)
        train_loss = F.cross_entropy(y_pred, data.y)
        total_loss += train_loss.item() * data.num_graphs
        train_loss.backward()
        optimizer.step()
    return total_loss / len(train_loader.dataset)

@torch.no_grad()
def test():
    model.eval()

    ys, preds = [], []
    for data in test_loader:
        ys.append(data.y)
        data = data.to(device)
        out, pred = model(data).max(dim=1)
        preds.append(pred.float().cpu())  

    y, pred = torch.cat(ys, dim=0).numpy(), torch.cat(preds, dim=0).numpy()
    correct=np.sum(y == pred)
    tot_ele = pred.size
    return f1_score(y, pred, average='micro') if pred.sum() > 0 else 0

for epoch in range(1, 101):
    loss = train()
    test_f1 = test()
    print('Epoch: {:02d}, Loss: {:.4f}, Test: {:.4f}'.format(
        epoch, loss, test_f1))

[rusty1s]

I modified your code as follows:

1. Added non-linearities:

x = self.conv1(x, edge_index).relu()
x = self.conv2(x, edge_index).relu()
x = self.conv3(x, edge_index)

2. The softmax and loss combination: cross_entropy already applies a softmax, so no need to do it in forward.
3. I increased the batch_size to 512.

Indeed, the model outputs all zero classes at the beginning, but accuracy improves later on (for me it was after around 40 epochs).

[JindeShubhamA]

Thank you for the reply. Yes, I am able to see improvement after 50+ epochs. Thanks will update you if I get some interesting results.