PyTorch Geometric

https://pyg.org/docs
Admin
PyTorch Geometric is a library for writing and training Graph Neural Networks (GNNs) on structured...
Tokens:1,590
Snippets:3
Trust Score:8.7
Update:1 year ago
Show doc for...
Tokens:
Raw
### PyG: Graph-level Prediction with TUDataset and GATConv

Source: https://pyg.org/docs/index

This example shows how to load a graph-level dataset (TUDataset) and train a GNN model using GATConv. It utilizes DataLoader for batching and includes a model with global pooling for graph-level outputs.

```python
# Load a dataset:
dataset = TUDataset(root, name="PROTEINS")

# Create a mini-batch loader:
loader = DataLoader(dataset, batch_size=256)

# Create your GNN model:
class GNN(torch.nn.Module):
    def __init__(self):
        self.conv1 = GATConv(dataset.num_features, 16)
        self.conv2 = GATConv(16, 16)
        self.lin = Linear(16, dataset.num_classes)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
         # Choose between different GNN building blocks:
        x = global_mean_pool(x, batch)
        return self.lin(x)

# Train you GNN model:
for data in loader:
    y_hat = model(data.x, data.edge_index, data.batch)
    loss = criterion(y_hat, data.y)
```

--------------------------------

### PyG: Node-level Prediction with Planetoid and GCNConv

Source: https://pyg.org/docs/index

This snippet demonstrates loading a node-level dataset (Planetoid) and setting up a GNN model using GCNConv for node classification. It includes creating a mini-batch loader and a basic training loop.

```python
# Load a dataset:
dataset = Planetoid(root, name="Cora")

# Create a mini-batch loader:
loader = NeighborLoader(dataset[0], num_neighbors=[25, 10])

# Create your GNN model:
class GNN(torch.nn.Module):
    def __init__(self):
         # Choose between different GNN building blocks:
        self.conv1 = GCNConv(dataset.num_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        return self.conv2(x, edge_index)

# Train you GNN model:
for data in loader:
    y_hat = model(data.x, data.edge_index)
    loss = criterion(y_hat, data.y)
```

--------------------------------

### PyG: Link-level Prediction with Reddit and GraphSAGE

Source: https://pyg.org/docs/index

This snippet illustrates link-level prediction tasks by loading the Reddit dataset and employing LinkNeighborLoader. It defines a GNN model with GraphSAGE for encoding node representations and an InnerProductDecoder for link prediction.

```python
# Load a dataset:
dataset = Reddit(root)

# Create a mini-batch loader:
loader = LinkNeighborLoader(dataset[0], num_neighbors=[25, 10])

# Create your GNN model:
class GNN(torch.nn.Module):
    def __init__(self):
         # Choose between different GNN building blocks:
        self.encoder = GraphSAGE(dataset.num_features, 16, num_layers=2)
        Self.decoder = InnerProductDecoder()

    def forward(self, x, edge_index, edge_label_index):
        x = self.encoder(x, edge_index)
        return self.decoder(x, edge_label_index)

# Train you GNN model:
for data in loader:
    y_hat = model(data.x, data.edge_index, data.edge_label_index)
    loss = criterion(y_hat, data.edge_label)
```
Tokens:
Raw