Perforated AI OGBN Example

[RorryB]

This PR is to include an example of how Perforated Backpropagation can be used to improve pytorch_geometric models. It adds on the previous PR which added a scheduler, both of which improve upon the test accuracy of the original ogbn_train.py example.

Perforated Backpropagation empowers each neuron in a PyTorch network with Dendrite nodes. These nodes are outside the network, but able to inform the neurons inside the network to make better decisions by leveraging a covariance based learning rule to reduce residual error. The PyTorch add-on combines modern neuroscience with modern software to create neural networks that are smarter, smaller, and more accurate.

Run docker from torch_geometric directory

docker run --gpus all -i --shm-size=8g -v .:/pai -w /pai -t nvcr.io/nvidia/pyg:24.11-py3 /bin/bash

Within Docker

pip install -e .
pip install safetensors
pip install perforatedai

Ran original with:

CUDA_VISIBLE_DEVICES=0 python examples/ogbn_train.py --dataset ogbn-products --batch_size 128 --model sage

Results:

Test Accuracy: 77.06%

Run PAI with:

PAIEMAIL=YourEmail PAITOKEN=YourToken CUDA_VISIBLE_DEVICES=0 python ogbn_train_perforatedai.py --dataset ogbn-products --batch_size 128 --saveName ogbnPAI --model sage

Results:

Test Accuracy: 78.05%

Update: Example updated to include base usage with default SGFormer.

Original:

CUDA_VISIBLE_DEVICES=0 python ogbn_train.py --dataset ogbn-products --batch_size 128

Results:

Test Accuracy: 80.86%

PAI:

    PAIEMAIL=YourEmail PAITOKEN=YourToken CUDA_VISIBLE_DEVICES=0 python ogbn_train_perforatedai.py --dataset ogbn-products --batch_size 128 --saveName ogbnPAI

Results:

Test Accuracy: 81.27%

[puririshi98]

LGTM but I would like @akihironitta to take a look since he may have a different opinion on how best to expose Perforated AI for PyG

[akihironitta]

This PR is to include an example of how Perforated Backpropagation can be used to improve pytorch_geometric models. It adds on the previous PR which added a scheduler, both of which improve upon the original results.

I'm trying to understand the change at a high level, and sorry I may not have enough context here. Could you elaborate more on the change in the PR description? Could you link the previous PR (and relevant issues if there is)? What is Perforated AI? What are the original results?

[RorryB]

This PR is to include an example of how Perforated Backpropagation can be used to improve pytorch_geometric models. It adds on the previous PR which added a scheduler, both of which improve upon the original results.

I'm trying to understand the change at a high level, and sorry I may not have enough context here. Could you elaborate more on the change in the PR description? Could you link the previous PR (and relevant issues if there is)? What is Perforated AI? What are the original results?

@akihironitta PR description has now been updated to include these details.

[RorryB]

Example has been updated to include usage with default SGFormer.

[akihironitta]

Sorry for the delay, and thanks for the example! Looks cool :)

I feel like this is a kind of example that should be maintained by https://github.com/PerforatedAI/PerforatedAI-API. Given our current bandwidth, I'm not confident that we can keep this example maintained and up-to-date in the future.

[puririshi98]

@RorryB i think the best way forward is to copy this example over to your public PAI repo, then make a reference to that in our README. so this PR would just be an update to PyG examples readme and the main example code would go in your own repo.

[RorryB]

That sounds great! Just updated the PR to only include changes to the README and CHANGELOG with a link to the Perforated AI examples repository.

[puririshi98]

lgtm