Graph Classification With PyG And W&B

Setup

We will start by installing the basic packages i.e. PyTorch Geometric for implementing our graph neural networks, plotly for easier visualization and W&B for tracking our experiments.

Now let us import wandb and log in to your W&B account.

Now we will import all the packages that we will use as we progress through this example.

Dataset and Exploratory Data Analysis

To start logging and tracking information with W&B, we will first initialize a run. This run will reside inside a project in your W&B profile and it will store everything related to this experiment which in this case is EDA.

Note: Usage of Weights & Biases for training is optional, in case you don't wish to use it, simply uncheck use_wandb.

We will work on the MUTAG dataset for classifying graphs into one of two classes. Downloading and loading the graphs and labels from this dataset is supported internally by PyG.

In the snippet below, we convert the PyG graph to a plotly figure for visualization.

In order to perform EDA, we will log the input graphs along with the number of nodes, edges and the ground truth labels as a W&B Table.

We would also like to log and track our dataset so we will add the downloaded dataset to a W&B artifact and log that as well.

Training the Model

Creating Training and Testing Data

We shuffle the dataset and split it into training and testing sets

Let us create the dataloaders to effectively batch the graph inputs. Batching graphs is implemented extremely effectively in PyG and that has been used in this example.

Implementing the Model

To perform the classification, we use a very simple model with three graph convolution layers implemented in PyTorch Geometric.

Training

In the following cell, start a new W&B run and using the use_artifact function, we tell W&B that the given artifact is being used as input to this run.

Note: Usage of Weights & Biases for training is optional, in case you don't wish to use it, simply uncheck use_wandb.

In the function in the cell below, we implement a function to calculate accuracy and loss on the test set but along with that we also create a W&B table which can be used to see the graph along with the output from the model and the ground truth label.

The great thing about PyG is that all the code is very similar to PyTorch and you can use the same training logic which is what has been done in the following cell. We calculate the accuracy and loss on the training and validation sets and long them to our W&B dashboard.

We also log the table that was mentioned above to the run dashboard for debugging.

This is what the project dashboard looks like and you can interact with it i.e. make the panels larger or smaller and more!

📚 Resources

• Node Classification (with W&B) contains a few tips for taking most advantage of W&B with PyG
• Point Cloud Classification using PyTorch Geometric
• Recommending Amazon Products using Graph Neural Networks in PyTorch Geometric

❓ Questions about W&B

If you have any questions about using W&B to track your model performance and predictions, please contact support@wandb.com