Graph Neural Networks with PyTorch

This repository contains small, readable implementations of Graph Neural Network building blocks and training experiments using PyTorch, PyTorch Geometric, and PyTorch Lightning.

The project currently focuses on:

• A from-scratch Graph Convolutional Network layer example.
• A from-scratch graph multi-head attention layer example.
• Node classification on the Cora citation network.
• Early graph classification work on the MUTAG dataset.

Project Structure

File	Description
gcn.py	Minimal GCN layer implemented with dense adjacency matrices.
gma.py	Graph multi-head attention layer implemented from scratch.
train.py	Node-level classification on the Cora dataset using MLP and GNN models.
gcn_classification.py	Graph-level classification experiment for the MUTAG dataset.
checkpoint.py	Seed setup, data/checkpoint paths, and pretrained checkpoint download helper.
module.py	Common imports and plotting/deep-learning setup utilities.
setup.py	Package metadata and dependency list.

Requirements

• Python 3.11 or newer
• PyTorch
• PyTorch Geometric
• PyTorch Lightning
• torchvision
• NumPy
• matplotlib
• seaborn
• tqdm

Install the package and dependencies in editable mode:

pip install -e .

If PyTorch Geometric installation fails, install the correct PyTorch Geometric wheels for your PyTorch and CUDA version from the official installation guide:

pip install torch_geometric

Data and Checkpoints

Datasets are downloaded automatically by PyTorch Geometric into:

./data_path

Model checkpoints are stored in:

./checkpoint_path

To download the pretrained checkpoint files referenced by the project, run:

python checkpoint.py

Running the Examples

1. Basic GCN Layer

Run the dense-adjacency GCN example:

python gcn.py

This script creates a small toy graph, applies a graph convolution layer, and prints the input features, adjacency matrix, and output features.

2. Graph Multi-Head Attention Layer

Run the graph attention example:

python gma.py

This script builds a small graph, applies the custom multi-head attention layer, and prints the resulting node embeddings.

3. Node Classification on Cora

Run the Cora node-classification training script:

python train.py

The script trains and evaluates:

• An MLP baseline that ignores graph edges.
• A GNN model using PyTorch Geometric message-passing layers.

It reports train, validation, and test accuracy.

Models

GNN

Defined in train.py, this class builds a configurable stack of PyTorch Geometric graph layers. Supported layer names include:

• GCN
• GAT
• Graph
• GraphConv

MLP

Defined in train.py, this baseline uses linear layers, ReLU activations, and dropout. It is useful for comparing graph-aware learning against a model that only uses node features.

GCN

The Lightning module in train.py wraps either the GNN or MLP model and handles training, validation, testing, optimizer setup, and metric logging.

Notes

gcn_classification.py contains an in-progress graph-level classification experiment for MUTAG. It shows the intended architecture for graph pooling and classification, but it may need cleanup before being used as a stable training entry point.