Implementing a GNN using the PyTorch Geometric library

In this section, we will implement a GNN using the PyTorch Geometric library, which simplifies the process of training GNNs. We apply the GNN to QM9, a dataset consisting of small molecules, to predict isotropic polarizability, which is a measure of a molecule’s tendency to have its charge distorted by an electric field.

Installing PyTorch Geometric

PyTorch Geometric can be installed via conda or pip. We recommend you visit the official documentation website at https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html to select the installation command recommended for your operating system. For this chapter, we used pip to install version 2.0.2 along with its torch-scatter and torch-sparse dependencies:

pip install torch-scatter==2.0.9
pip install torch-sparse==0.6.12
pip install torch-geometric==2.0.2

Let’s start by loading a dataset of small molecules and look at...