The difference between NeighborSampler and NeighborLoader during training of the model

[Baizx98]

I'm currently studying GraphSAGE, and I've noticed that there are two types of samplers for sampling: NeighborSampler and NeighborLoader. However, they return iterators with different data types. The way I pass parameters for each SAGE layer in the model's forward function is also different, and there are differences in the train function as well. I'm not sure what the difference is between these two approaches.

# use NeighborSampler
import torch
from torch_geometric.data import NeighborSampler
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv

class SAGE(torch.nn.Module):
   def __init__(self, in_channels, hidden_channels, out_channels):
       super(SAGE, self).__init__()
       self.conv1 = SAGEConv(in_channels, hidden_channels)
       self.conv2 = SAGEConv(hidden_channels, out_channels)

   def forward(self, x, data_flow):
      # Different convolutional layers receive different adjacency matrices
       x = self.conv1(x, data_flow[0])
       x = x.relu()
       x = self.conv2(x, data_flow[1])
       return x.log_softmax(dim=-1)

data = OGB_MAG(root='path/to/dataset', name='ogbn-products')[0]

train_loader = NeighborSampler(data, size=[10, 10], num_hops=2, batch_size=100, shuffle=True, add_self_loops=True)

model = SAGE(dataset.num_features, 16, dataset.num_classes)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

def train():
   model.train()
   total_loss = 0
   for batch_size, n_id, adjs in train_loader:
       adjs = [adj.to(device) for adj in adjs]
       optimizer.zero_grad()

       out = model(data.x[n_id.to(device)], adjs)
       loss = F.nll_loss(out, data.y[n_id[:batch_size]])
       loss.backward()
       optimizer.step()
       total_loss += loss.item()
   return total_loss / len(train_loader)

# use NeighborLoader
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
root = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'products')
dataset = PygNodePropPredDataset('ogbn-products', root)
split_idx = dataset.get_idx_split()
evaluator = Evaluator(name='ogbn-products')
data = dataset[0].to(device, 'x', 'y')

train_loader = NeighborLoader(
    data,
    input_nodes=split_idx['train'],
    num_neighbors=[15, 10, 5],
    batch_size=1024,
    shuffle=True,
    num_workers=12,
    persistent_workers=True,
)

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

        self.num_layers = num_layers

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

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

    def forward(self, x, edge_index):
        for i, conv in enumerate(self.convs):
            # Each convolutional layer uses the same edge_index.
            x = conv(x, edge_index)
            if i != self.num_layers - 1:
                x = x.relu()
                x = F.dropout(x, p=0.5, training=self.training)
        return x


model = SAGE(dataset.num_features, 256, dataset.num_classes, num_layers=3)
model = model.to(device)

def train(epoch):
    model.train()

    pbar = tqdm(total=split_idx['train'].size(0))
    pbar.set_description(f'Epoch {epoch:02d}')

    total_loss = total_correct = 0
    for batch in train_loader:
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index.to(device))[:batch.batch_size]
        y = batch.y[:batch.batch_size].squeeze()
        loss = F.cross_entropy(out, y)
        loss.backward()
        optimizer.step()

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

    pbar.close()

    loss = total_loss / len(train_loader)
    approx_acc = total_correct / split_idx['train'].size(0)

    return loss, approx_acc

Additionally, I'd like to know if the number of SAGE layers in a GNN model is the same as the depth K in the GraphSAGE algorithm. Does the number of SAGE layers represent the hop of aggregated neighbors? How should I understand the num_neighbors in the NeighborLoader? How are these depth, layer, and hop related?

[LukeLIN-web]

NeighborSampler was deprecated so that you could use NeighborLoader.
Yes, the number of SAGE layers in a GNN model is the same as the depth K in the GraphSAGE algorithm.
num_neighbors [25,10] means, first hop, the target node will sample 25 neighbors of it, and then these nodes(<=26) will sample 10 neighbors of them. So one mini-batch ( aka one subgraph) will have <=260 nodes.

[Baizx98]

Thanks！I got it！