SGC in Pytorch Geometric vs original SGC speed

[mdanb]

I tested the same dataset on a slightly modified version of Pytorch Geometric's SGC, as well as on a slightly modified version of the official SGC implementation. The original SGC is much faster. Also, it is much more stable. I'm wondering why that might be.

Some (maybe) relevant details:
My task is graph classification, where the graph is static. There are 22113 nodes, each with 1 feature, and there are 3 labels for the graphs. I use SparseTensor for my edges in Pytorch geometric for memory efficiency.
In both cases, I concatenate resulting node embeddings after message passing and use that for classification (which is not an issue since the graph is static)

Here are the two versions of the code:
Pytorch Geometric:

class MySGConv(MessagePassing):
    _cached_x: Optional[Tensor]

    def __init__(self, in_channels: int, out_channels: int, K: int = 1,
                 cached: bool = False, add_self_loops: bool = True,
                 bias: bool = True, **kwargs):
        kwargs.setdefault('aggr', 'add')
        super(MySGConv, self).__init__(**kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.K = K
        self.cached = cached
        self.add_self_loops = add_self_loops

        self._cached_x = None

        self.lin = Linear(22113, out_channels, bias=bias)

        self.reset_parameters()

    def reset_parameters(self):
        self.lin.reset_parameters()
        self._cached_x = None

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:
        """"""
        cache = self._cached_x
        if cache is None:
            if isinstance(edge_index, Tensor):
                edge_index, edge_weight = gcn_norm(  # yapf: disable
                    edge_index, edge_weight, x.size(self.node_dim), False,
                    self.add_self_loops, dtype=x.dtype)
            elif isinstance(edge_index, SparseTensor):
                edge_index = gcn_norm(  # yapf: disable
                    edge_index, edge_weight, x.size(self.node_dim), False,
                    self.add_self_loops, dtype=x.dtype)

            for k in range(self.K):
                x = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                                   size=None)
                if self.cached:
                    self._cached_x = x
        else:
            x = cache

        x = torch.squeeze(x)
        return self.lin(x)

    def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
        return matmul(adj_t, x, reduce=self.aggr)

    def __repr__(self):
        return '{}({}, {}, K={})'.format(self.__class__.__name__,
                                         self.in_channels, self.out_channels,
                                         self.K)

class GNN(torch.nn.Module):
    def __init__(self, input_dim, output_dim, config, edge_attr, edge_index, device):
        super(GNN, self).__init__()
        self.num_layers = config['num_layers']
        self.convs = torch.nn.ModuleList()
        self.conv_type = config['conv_type']
        self.edge_attr = torch.Tensor(edge_attr).float()
        self.edge_index_sp = SparseTensor(row=edge_index[0], col=edge_index[1], value=self.edge_attr, sparse_sizes=(22113, 22113)).cuda()
        self.device = device
        self.convs.append(self.build_conv_model(input_dim, self.conv_type))

    def build_conv_model(self, input_dim, conv_type):
        if (conv_type == 'mySGC'):
            return MySGConv(in_channels=input_dim, out_channels=3,K=1)

    def forward(self, data):
        x, batch, edge_index, edge_attr = data.x, data.batch, self.edge_index, self.edge_attr
        batch_size = data.y.shape[0]
        x = x.reshape(batch_size, 22113, 1)
        for i in range(self.num_layers):
            x = self.convs[i](x, self.edge_index_sp)
        return F.log_softmax(x, dim=1)

    def loss(self, pred, label):
        return F.nll_loss(pred, label)

    def weight_reset(self, m):
        if isinstance(m, torch.nn.Linear):
            m.reset_parameters()

Original SGC:
Function used for message passing, prior to feature transformation through SGC:

def sgc_precompute(features, adj, degree=1):
    features = torch.transpose(features, 1, 0)
    for i in range(degree):
        features = torch.spmm(adj, features)
    return torch.transpose(features, 1, 0)

class SGC(nn.Module):
    """
    A Simple PyTorch Implementation of Logistic Regression.
    Assuming the features have been preprocessed with k-step graph propagation.
    """
    def __init__(self, nfeat, nclass):
        super(SGC, self).__init__()
        nclass = nclass
        self.post_mp = torch.nn.Linear(nfeat, nclass)
        self.init()

    def init(self):
        stdv = 1. / math.sqrt(self.post_mp.weight.size(1))
        self.post_mp.weight.data.uniform_(-stdv, stdv)

    def weight_reset(self, m):
        if isinstance(m, nn.Linear):
            m.reset_parameters()

    def forward(self, x):
        return F.log_softmax(self.post_mp(x))

[rusty1s]

The reason why it is faster is due to the pre-computation that is happening in SGC. We currently provide such an option via the cached=True argument in SGConv, but which expects working on a single graph.

Note that you can also just pre-compute the features in PyG, e.g., via the SIGN transform and input those into a single linear layer.

The only reason I can think of why it might be more stable is due to the different initialization techniques of the linear layers.

[mdanb]

When you say "single graph", what do you mean by that? That the graph edges are constant and node features can vary?

[rusty1s]

I mean learning on a single graph in a semi-supervised fashion, e.g., on Cora.