NeighborLoader vs NeighborSampler: Which should be used for GraphSAGE?

[bwdeng20]

Though NeighborSampler will be deprecated, I think it's layer-by-layer sampling results seem more close to the original GraphSAGE, a typical implmentation of which is like ogbn_products_sage.py under the examples directory.

class SAGE1(torch.nn.Module): 
#from https://github.com/pyg-team/pytorch_geometric/blob/d49735358258b3613a0463e131fac7c4e3198b8a/examples/ogbn_products_sage.py
    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)

Given a sampled cora subgraph rooted at 1000 (with torch.manual_seed(0)):

SAGE1 works as the original paper where each layer only propagates wihtin a part of the subgraph:

1st prop: [298,327] --> [1325], [746,1347] --> [281]
2nd prop: [1325,281] --> [1000]

But SAGE2 with NeighborLoader in examples/reddit.py seems to have a different aggregation process where each layer propagates messages wihtin the entire subgraph:

1st prop: [298,327] --> [1325], [746,1347] --> [281],[1325,281] --> [1000]
2nd prop: [298,327] --> [1325], [746,1347] --> [281],[1325,281] --> [1000]

class SAGE2(torch.nn.Module):
#from https://github.com/pyg-team/pytorch_geometric/blob/master/examples/reddit.py
    def forward(self, x, edge_index):
        for i, conv in enumerate(self.convs):
            x = conv(x, edge_index)
            if i < len(self.convs) - 1:
                x = x.relu_()
                x = F.dropout(x, p=0.5, training=self.training)
        return x

Overally, my questions are:

• Are the above statements correct?
• Does the recommended NeighborLoader lead to a SAGE different from the paper (and SAGE1 in the first code snippet)?
• Did you find that SAGE2 in the second snippet with NeighborLoader have a better or on-par peformance with that ? So you prefer a NeighborLoader-style way.

[rusty1s]

Yes, this is correct. Please feel free to continue using the NeighborSampler if it fits you better. Overall, I don't like it that much because it leads to complicated GNN designs that are tightly coupled to a specific sampler.

The NeighborLoader samples neighbors equivalently, but will group them into a single edge_index matrix. Note that the edge_index is still directed though. As a result, the output is the same. The only difference is that NeighborLoader will be marginally slower since we are computing node embeddings for nodes we no longer need. This is a trade-off we make to obtain a cleaner GNN design.

[bwdeng20]

Thanks for the clarification! For SAGEConv in PyG, I think these two styles coincide with each other only if the sampled edge_index contains no self-loops and that's indeed the case of GraphSAGE. Is this correct?

[rusty1s]

Even in case of self-loops it shouldn't matter. Can you clarify why it should?

[bwdeng20]

Let's say the subgraph in my original post is a self-looped graph. We immediately update the computational graphs of SAGE1 and SAGE2 as:

1st prop: [298,327,1325] --> [1325], [746,1347,281] --> [281]
2nd prop: [1325,281,1000] --> [1000]

and

1st prop: [298,327,1325] --> [1325], [746,1347,281] --> [281],[1325,281,1000] --> [1000]
2nd prop: [298,327,1325] --> [1325], [746,1347,281] --> [281],[1325,281,1000] --> [1000]

Now for the target node 1000 , SAGE2 aggregates information from the neighbor twice while SAGE1 still performs one time convolution.

[rusty1s]

Please note that in the first prop, [1325,281,1000] --> [1000] is also included in NeighborSampler to allow for the application of skip-connections.

[bwdeng20]

Thanks!

[Baizx98]

Hey, I have the same question. Does this mean that the two aggregation methods used in SAGE1 and SAGE2 will eventually yield the same results? If so, how should I decide the number of layers for sampling and the number of convolutional layers in the model?

[LukeLIN-web]

the embedding of node id 1000 is the same. The only difference is that NeighborLoader will be marginally slower since we are computing node embeddings for nodes we no longer need( node id 1325,281 ).

[Baizx98]

Thanks！我悟了😊