Training on an entire graph at once can strain memory. A graph attention network (GAT), for example, calculates attention scores for every head and edge, while autograd retains tensors needed for training. Mini-batching loads smaller portions of the graph into the training loop instead of using the whole graph as one batch. section 4.4.1
Batching and sampling are related but distinct. Batching divides training data into smaller groups of nodes or edges; sampling chooses what goes into each group. A group may form a connected subgraph, but need not. Processing one group at a time reduces the need to hold the entire graph in memory during an epoch. section 7.7.1
A practical node-batching workflow is:
- Choose the training nodes and a batch size. In PyTorch Geometric, NeighborLoader can use input_nodes=data.train_mask and batch_size=128. section 4.4.1
- Set num_neighbors for each GNN layer to control neighborhood sampling. The GAT example uses [1000]*2 for two layers. section 4.4.1
- Iterate over the loader in the training loop, move each sampled batch to the device, then run the model, loss calculation, backward pass, and optimizer step. section 4.4.1
Listing 4.6 Setting up NeighborLoader for GAT
from torch_geometric.loader import NeighborLoader
batch_size = 128
loader = NeighborLoader(
data,
num_neighbors=[1000]*2, #1
batch_size=batch_size, #2
input_nodes=data.train_mask)
sampled_data = next(iter(loader))
print(f'Checking that batch size is \
{batch_size}: {batch_size == \
sampled_data.batch_size}')
print(f'Percentage fraud in batch: \
{100*sampled_data.y.sum()/\
len(sampled_data.y):.4f}%')
sampled_data
The NeighborLoader setup shows the three inputs that determine this example’s batches: training nodes, batch size, and neighbors sampled across layers. section 4.4.1
Choosing batches for a growing geographic graph
For a graph of geographic areas connected by proximity, roads, or development patterns, batch size and sampled-neighbor count offer tuning controls as the graph grows. GraphSAGE uses neighborhood sampling and mini-batching to limit its dependence on the full graph size; a GCN can also use PyG mini-batches rather than store the entire adjacency matrix. section 7.6.1
Smaller batches are not a free substitute for full-graph training. Sampling can omit neighbors needed for message passing or introduce bias, affecting performance and generalizability. Samplers also differ in epoch and convergence time, so test candidates on limited data. Batching can be combined with other scale measures, including sparse representations, parallel computing, and remote backends. section 7.7.1 section 7.11
