Skip to main content

Graph Data Augmentations for PyTorch Geometric

Project description


Graph Data Augmentation Library for PyTorch Geometric.

What is it?

Data augmentations are heavily used in Computer Vision and Natural Language Processing to address data imbalance, data scarcity, and prevent models from overfitting. They have also proven to yield good results in both supervised and self-supervised (contrastive) settings.

grafog (portmanteau of "graph" and "augmentation") provides a set of methods to perform data augmentation on graph-structured data, especially meant for self-supervised node classification. It is built on top of torch_geometric and is easily integrable with its Data API.


You can install the library via pip:

$ pip install grafog

You can also install the library from source:

$ git clone
$ cd grafog
$ pip install -e .




The library comes with the following data augmentations:

  1. Random Node Drop
  2. Random Edge Drop
  3. Normalize Features
  4. MixUp Strategy
  5. Node Feature Masking
  6. Edge Feature Masking

There are many more features to be added over time, so stay tuned!

from torch_geometric.datasets import CoraFull
import grafog.transforms as T

node_aug = T.Compose([
    T.NodeMixUp(lamb=0.5, classes=7),

edge_aug = T.Compose([

data = CoraFull()
model = ...

for epoch in range(10): # begin training loop
    new_data = node_aug(data) # apply the node augmentation(s)
    new_data = edge_aug(new_data) # apply the edge augmentation(s)
    x, y = new_data.x, new_data.y


This library was built as a project for a class (UIT2201) at NUS. I planned and built it over the span of 10 weeks. I thank Prof. Mikhail Filippov for his guidance, feedback, and support!

If you spot any issues, feel free to raise a PR or Issue. All meaningful contributions welcome!



Project details

Release history Release notifications | RSS feed

This version


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

grafog-0.1.tar.gz (3.8 kB view hashes)

Uploaded source

Built Distribution

grafog-0.1-py3-none-any.whl (4.5 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page