PyTorch Geometric (PyG) is a Python library for building graph neural networks with PyTorch. Use it when relationships between entities matter to a prediction: a node’s representation can depend on its neighbors, not just its own features. Before choosing a GNN, decide whether the data has useful graph structure and whether a simpler model is a better fit. section 1.1 section 4.3.2
A practical PyG workflow moves from raw records to graph tensors and then to training:
- Choose nodes, edges, and their features. Raw records may first become edge lists or adjacency lists and a NetworkX graph. section 1.5.1 section 8.0
- Convert the graph with `from_networkx`, then assign tensors such as `data.x` for node features and `data.y` for labels. Prepare training and test indices or masks without losing the graph structure. section 2.2.1 section 4.3.2
- Define PyTorch modules using graph layers such as `GCNConv` or `SAGEConv`. Pass node features and `edge_index` through the layers to produce embeddings or predictions. section 2.2.1 section 3.2.2
Node classification with a GCN
A two-layer `GCNConv` model reads `data.x` and `data.edge_index`, applies an activation and dropout between layers, and produces node predictions. The training mask selects nodes for the loss, but the forward pass receives the whole graph so message passing can use neighboring nodes. If features are absent, initialized node features can provide a starting point for learned embeddings. section 4.3.2 section 2.2.1
Listing 2.4 Our SimpleGNN class
class SimpleGNN_embeddings(torch.nn.Module):
def __init__(self, num_features, hidden_channels): #1
super(SimpleGNN, self).__init__()
self.conv1 = GCNConv(num_features, \
hidden_channels) #2
self.conv2 = GCNConv(hidden_channels,\
hidden_channels) #3
def forward(self, x, edge_index): #4
x = self.conv1(x, edge_index) #5
x = torch.relu(x) #6
x = torch.dropout(x, p=0.5, train=self.training) #7
x = self.conv2(x, edge_index) #8
return x #9
A compact PyTorch module shows how two PyG graph-convolution layers consume `x` and `edge_index` to produce node embeddings. section 2.2.1
PyG also lets you vary neighborhood aggregation: GraphSAGE layers can use different methods across layers or combine methods such as mean, max, and sum within one layer. Test such choices rather than assuming added complexity helps. Small graphs can fit in laptop or desktop memory; data beyond one machine’s capacity calls for separate handling strategies. section 3.2.2 section 1.5.1
