dmlc / dmlc/dgl

[DataLoader] Support weighted seed node/edge sampling in DGL DataLoader

Open
#3,431 4 comments 0 reactions 2 assignees Claimed by @BarclayII View on GitHub
feature request topic: Dataloader
Dominant language
Python
Stars
14.3k
Forks
3.1k
PR merge metrics
No merged PRs in 30d

Description

## Adding Torch Sampler with NodeDataLoader

Attempting Node Classification with GraphSAGE and Minibatching for imbalanced classes does not really work.
When trying to add a Pytorch `WeightedRandomSampler` together with dgl `NodeDataLoader` and `MultiLayerFullNeighborSampler` this results to a weird error about missing features (on the fourth iteration):
```sh
dgl._ffi.base.DGLError: Expected data to have 32 rows, got 30.
```

## Script for reproducing the error on the Cora Dataset

Requirements are :
- dgl
- numpy
- torch
- argparse
- tqdm

```python
import torch
torch.manual_seed(0)
import random
random.seed(0)
import numpy as np
np.random.seed(0)

import argparse

import dgl.dataloading

import torch.nn as nn
from tqdm import tqdm
from dgl.nn import SAGEConv

class MyModel(nn.Module):
def __init__(self, in_feats, num_classes):
super(MyModel, self).__init__()
self.conv = SAGEConv(in_feats, num_classes, aggregator_type="gcn")

def forward(self, mfgs, x):
# Lines that are changed are marked with an arrow: "<---"
h = self.conv(mfgs[0], x) # <--- # <---
return h

def get_sample_weights(labels):
class_sample_count = torch.tensor([(labels == t).sum() for t in torch.unique(labels, sorted=True)])
weight = 1. / class_sample_count.float()
sample_weights = torch.tensor([weight[t] for t in labels])
return sample_weights

def main(config):
"""Pass parameters to create experiment"""

# --------------- Dataset Loading -------------------------
dataset = dgl.data.CoraGraphDataset()
g = dataset[0]
g = dgl.add_self_loop(g)
n_classes = dataset.num_classes
train_nid = torch.tensor(range(g.num_nodes())).type(torch.int64)
in_feats = g.ndata["feat"].shape[1]

# --------------- Transfer to Devise ---------------
use_cuda = config["gpu"] >= 0 and torch.cuda.is_available()
if use_cuda:
device = torch.device('cuda:%d' % config["gpu"])
dataloader_device = device
else:
device = torch.device('cpu')
train_nid = train_nid.to(device)
dataloader_device = device
# ---------------- Sampler Definition ---------------
# Graph Sampler takes all available neighbors
graph_sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)

# Balance Label Sampler
label_weights = get_sample_weights(g.ndata["label"])
# Torch Sampler
sampler = torch.utils.data.sampler.WeightedRandomSampler(label_weights, len(label_weights))

dataloader = dgl.dataloading.NodeDataLoader(
g,
train_nid,
graph_sampler,
device=dataloader_device,
batch_size=config["batch_size"],
drop_last=False,
num_workers=0,
sampler=sampler
)

# Define model and optimizer
model = MyModel(in_feats, n_classes)
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# Training loop
for epoch in tqdm(range(config["num_epochs"]), desc='epoch'):
# Loop over the dataloader to sample the computation dependency graph as a list of blocks.
for step, (input_nodes, seeds, blocks) in enumerate(tqdm(dataloader, position=0, leave=True, desc='data')):
# Load the input features as well as output labels
batch_inputs = blocks[0].srcdata["feat"]
batch_labels = blocks[-1].dstdata['label']
print(step, blocks[0].num_nodes(), batch_labels.shape)
# Predict and loss
batch_pred = model(blocks, batch_inputs)
loss = criterion(batch_pred, batch_labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()

if __name__ == '__main__':
argparser = argparse.ArgumentParser(description='Weighted Sampling SAGE')
argparser.add_argument('--gpu', type=int, default=-1,
help="GPU device ID. Use -1 for CPU training")
argparser.add_argument('-d', '--dataset', type=str, default='toy01')
argparser.add_argument('--num-epochs', type=int, default=20)
argparser.add_argument('--batch-size', type=int, default=32)

args = argparser.parse_args()
config = vars(args)

main(config)
```

## The Error Log
```sh
Traceback (most recent call last):
File "Path/to/Script/test_weighted_sampler.py", line 109, in
main(config)
File "Path/to/Script/test_weighted_sampler.py", line 91, in main
batch_pred = model(blocks, batch_inputs)
File "/Path/To/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
return forward_call(*input, **kwargs)
File "Path/to/Script/test_weighted_sampler.py", line 25, in forward
h = self.conv(mfgs[0], x) # <--- # <---
File "/Path/To/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
return forward_call(*input, **kwargs)
File "/Path/To/lib/python3.9/site-packages/dgl/nn/pytorch/conv/sageconv.py", line 243, in forward
graph.dstdata['h'] = graph.srcdata['h'][:graph.num_dst_nodes()]
File "/Path/To/lib/python3.9/site-packages/dgl/view.py", line 81, in __setitem__
self._graph._set_n_repr(self._ntid, self._nodes, {key : val})
File "/Path/To/lib/python3.9/site-packages/dgl/heterograph.py", line 4118, in _set_n_repr
self._node_frames[ntid].update(data)
File "/Path/To/lib/python3.9/_collections_abc.py", line 940, in update
self[key] = other[key]
File "/Path/To/lib/python3.9/site-packages/dgl/frame.py", line 405, in __setitem__
self.update_column(name, data)
File "/Path/To/lib/python3.9/site-packages/dgl/frame.py", line 478, in update_column
raise DGLError('Expected data to have %d rows, got %d.' %
dgl._ffi.base.DGLError: Expected data to have 32 rows, got 30.

Process finished with exit code 1
```
Thank you for your help!!!

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.