PyTorch Geometric Setup on DGX

README.md

PyTorch Geometric Setup on DGX

Steps

1. Create a Python virtual enviroment and activate it

[user@hal-dgx]$ pyenv virtualenv 3.9.6 pytorch-geometric-example
[user@hal-dgx]$ pyenv activate pytorch-geometric-example
(pytorch-geometric-example) [user@hal-dgx]$ python -m pip install --upgrade pip setuptools wheel

2. Install the requirements

(pytorch-geometric-example) [user@hal-dgx]$ bash install.sh
...
# python -m pip list | grep 'torch'
torch                 1.9.1+cu111
torch-geometric       2.0.1
torch-scatter         2.0.8
torch-sparse          0.6.12

# Version checks
Torch version: 1.9.1+cu111
Torch detected CUDA version: 11.1
# nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Wed_Jul_22_19:09:09_PDT_2020
Cuda compilation tools, release 11.0, V11.0.221
Build cuda_11.0_bu.TC445_37.28845127_0

3. Run the example (for 20 epcohs as sanity check)

(pytorch-geometric-example) [user@hal-dgx]$ python example.py
Accuracy: 0.7890

example.py

import torch
import torch.nn.functional as F
from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv


class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_node_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index

        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, training=self.training)
        x = self.conv2(x, edge_index)

        return F.log_softmax(x, dim=1)


if __name__ == "__main__":
    # c.f. https://pytorch-geometric.readthedocs.io/en/latest/notes/introduction.html#learning-methods-on-graphs
    dataset = Planetoid(root="/tmp/Cora", name="Cora")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = GCN().to(device)
    data = dataset[0].to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

    model.train()
    # for epoch in range(200):
    for epoch in range(20):
        optimizer.zero_grad()
        out = model(data)
        loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
        loss.backward()
        optimizer.step()

    model.eval()
    pred = model(data).argmax(dim=1)
    correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
    acc = int(correct) / int(data.test_mask.sum())
    print("Accuracy: {:.4f}".format(acc))

install.sh

#!/bin/bash

printf "\n# Install dependnecies\n$ python -m pip install -r requirements.txt\n"
python -m pip install -r requirements.txt

printf "\n# python -m pip list | grep 'torch'\n"
python -m pip list | grep 'torch'

printf "\n# Version checks\n"
python -c 'import torch; print(f"Torch version: {torch.__version__}")'
python -c 'import torch; print(f"Torch detected CUDA version: {torch.version.cuda}")'
printf "# nvcc --version\n"
nvcc --version

requirements.txt

--find-links https://download.pytorch.org/whl/torch_stable.html
torch==1.9.1+cu111
--find-links https://data.pyg.org/whl/torch-1.9.1+cu111.html
# Let patch releases float for bug fixes
torch-scatter~=2.0.8
torch-sparse~=0.6.12
torch-geometric~=2.0.1