SGConv for ogbn-products dataset

[MohsenFazaeli]

Hi,

I am going to run SGC on ogbn-products dataset. In PYG implementation of SGC S(S^{n-1}X) W is calculated which propagates data through normalized adjacency matrix n iteration which does not fit in memory. Can we calculate this higher powers of S fist S^{n}X W and then use batching to decrease memory requirements?

[EdisonLeeeee]

Yes, you can precompute $S$ manually outside the network, followed by an MLP with minibatch training.

[rusty1s]

Got it. You can also take a look at the SIGN transform, which basically does what you want.

[EdisonLeeeee]

@MohsenFazaeli Following @rusty1s's suggestion, this can be simply implemented via:

import torch

from torch_geometric.data import Data
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.transforms import BaseTransform


@functional_transform('sgc')
class SGC(BaseTransform):
    def __init__(self, K: int):
        self.K = K

    def __call__(self, data: Data) -> Data:
        from torch_sparse import SparseTensor

        assert data.edge_index is not None
        row, col = data.edge_index
        adj_t = SparseTensor(row=col, col=row,
                             sparse_sizes=(data.num_nodes, data.num_nodes))

        deg = adj_t.sum(dim=1).to(torch.float)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        adj_t = deg_inv_sqrt.view(-1, 1) * adj_t * deg_inv_sqrt.view(1, -1)

        assert data.x is not None
        x = data.x
        for i in range(1, self.K + 1):
            x = adj_t @ x
        data.x = x

        return data

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(K={self.K})'


from torch.utils.data import DataLoader
data = SGC()(data)
model = nn.Sequential(nn.Linear(...), ...)
for x in DataLoader(data.x, batch_size=128):
    pred = model(x)

[MohsenFazaeli]

@EdisonLeeeee , thank you for your detailed answer. I completed the code. Would you take a look at it, please?

import torch
import torch.nn as nn
from torch_geometric.data import Data
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.transforms import BaseTransform
from ogb.nodeproppred import Evaluator, PygNodePropPredDataset

import argparse
import os.path as osp

@functional_transform('sgc')
class SGC(BaseTransform):
    def __init__(self, K: int):
        self.K = K

    def __call__(self, data: Data) -> Data:
        from torch_sparse import SparseTensor

        assert data.edge_index is not None
        row, col = data.edge_index
        adj_t = SparseTensor(row=col, col=row,
                             sparse_sizes=(data.num_nodes, data.num_nodes))

        deg = adj_t.sum(dim=1).to(torch.float)
        deg_inv_sqrt = deg.pow(-0.5)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        adj_t = deg_inv_sqrt.view(-1, 1) * adj_t * deg_inv_sqrt.view(1, -1)

        assert data.x is not None
        x = data.x
        for i in range(1, self.K + 1):
            x = adj_t @ x
        data.x = x

        return data

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}(K={self.K})'

def main(args):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'OGB')
    #dataset = PygNodePropPredDataset(args.dataset, root, transform=T.ToSparseTensor())
    dataset = PygNodePropPredDataset(args.dataset, root)
    data = dataset[0]
    x = data.x.to(device)
    y = data.y.squeeze().to(device)
    split_idx = dataset.get_idx_split()
    train_idx = split_idx['train']
    val_idx = split_idx['valid']
    test_idx = split_idx['test']
    from torch.utils.data import DataLoader
    data = SGC(K=2)(data)
    model = nn.Sequential(
          nn.Linear(dataset.num_features, 64),
          nn.ReLU(),
          nn.Linear(64, dataset.num_classes)
        ).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    train_dataset = torch.utils.data.TensorDataset(x[train_idx], y[train_idx])
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=1024, shuffle=True)

    val_dataset = torch.utils.data.TensorDataset(x[val_idx], y[val_idx])
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1024, shuffle=False)

    test_dataset = torch.utils.data.TensorDataset(x[test_idx], y[test_idx])
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1024, shuffle=False)

    best_val_acc = final_test_acc = 0
    for epoch in range(1, 201):
        model.train()
        total_loss = total_correct = 0
        for x_batch,y_batch in train_loader:
            x_batch,y_batch = x_batch.to(device),y_batch.to(device)
            optimizer.zero_grad()
            out = model(x_batch)
            loss = nn.CrossEntropyLoss()(out, y_batch)
            loss.backward()
            optimizer.step()
            total_loss += float(loss)
            total_correct += int(out.argmax(dim=-1).eq(y_batch).sum())

        loss = total_loss / len(train_loader)
        approx_acc = total_correct / train_idx.size(0)
        print(f'Epoch {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {approx_acc:.4f}')

        model.eval()
        total_correct = 0
        for x_batch, y_batch in train_loader:
            x_batch, y_batch = x_batch.to(device), y_batch.to(device)
            out = model(x_batch)
            total_correct += int(out.argmax(dim=-1).eq(y_batch).sum())
        train_acc = total_correct / train_idx.size(0)

        total_correct = 0
        for x_batch, y_batch in val_loader:
            x_batch, y_batch = x_batch.to(device), y_batch.to(device)
            out = model(x_batch)
            total_correct += int(out.argmax(dim=-1).eq(y_batch).sum())
        val_acc = total_correct / val_idx.size(0)

        total_correct = 0
        for x_batch, y_batch in test_loader:
            out = model(x_batch)
            total_correct += int(out.argmax(dim=-1).eq(y_batch).sum())
        test_acc = total_correct / test_idx.size(0)

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            final_test_acc = test_acc
        print(f'Train: {train_acc:.4f}, Val: {val_acc:.4f}, '
              f'Test: {test_acc:.4f} FinalTest: {final_test_acc:.4f}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='OGBN-Products (SIGN)')
    parser.add_argument('--K', type=int, default=2)#SGC
    parser.add_argument('--dataset', type=str, default='ogbn-products')#SGC
    args = parser.parse_args()
    main(args)

[EdisonLeeeee]

Almost looks good to me. One major problem is that you need to transform the feature first:

...
data = SGC(K=2)(data)
x = data.x.to(device)
...

Or you can simply do it via:

dataset = PygNodePropPredDataset(args.dataset, root, transform=SGC(K=2))

And here are some minor suggestions:

• The original SGC paper only uses one single linear layer as the learning model. Not sure if a deep model would benefit the performance.
• SGC is sensitive to lr and weight_decay, you need to carefully tune them to achieve better results.
• If you find the performance is poor, adding BatchNorm1d layer might be a good choice.

Hope it helps!

[MohsenFazaeli]

Thank you! they are very helpful.