How to implement the more advanced sampler (e.g., ClassStratifiedSampler)?

[ha-lins]

Hi there,

Thanks for the powerful PyG.

• Background

I have a need of implementing the API of sampler (i.e., the ClassStratifiedSampler here).

Concretely, I'd like to iter some TUDataset graphs (like images) with different classes within a batch. For example, a batch can contain graphs with labels [0, 0, 1, 1, 0, 0, 1, 1].

• Problems

I have implemented the ClassStratifiedSampler and it returns target_indices of graphs correctly. For example, in the MUTAG of TUDataset, ClassStratifiedSampler can get desired graph index [169, 155, 166, 187, 20, 118, 91, 6, 108, 2, 173, 131, 81, 146, 60, 129] in a batch. However, the StopIteration issue occur as following. The code part is:

from torch_geometric.loader import DataLoader

supervised_sampler = ClassStratifiedSampler(data_source=dataset, world_size=1,
                                                rank=0,
                                                batch_size=8,
                                                classes_per_batch=2,
                                                seed=0,
                                                unique_classes=False)
supervised_loader = DataLoader(dataset,
                               sampler=supervised_sampler,
                               batch_size=8,
                               drop_last=True,
                               pin_memory=True,
                               num_workers=1)
batch = next(iter(supervised_loader))
print(batch)

I refer to the code here and it is originally written as:

        dist_sampler = ClassStratifiedSampler(
            data_source=dataset,
            world_size=world_size,
            rank=rank,
            batch_size=batch_size,
            classes_per_batch=classes_per_batch,
            seed=_GLOBAL_SEED,
            unique_classes=unique_classes)
        data_loader = torch.utils.data.DataLoader(
            dataset,
            batch_sampler=dist_sampler,
            pin_memory=True,
            num_workers=8)

The bug reports:

Traceback (most recent call last):
  File "Supervised_Dataset.py", line 344, in <module>
    batch = next(iter(supervised_loader))
  File "/home/liuchen/anaconda3/envs/semiGraph/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 435, in __next__
    data = self._next_data()
  File "/home/liuchen/anaconda3/envs/semiGraph/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1058, in _next_data
    raise StopIteration
StopIteration

It is reported in the source code here. I noticed there have been already some samplers (e.g., RandomIndexSampler), while they don't seem to match my need.

Questions

• Since the torch_geometric.loader.DataLoader inherits basically from torch.utils.data.DataLoader, it can be specified the sampler=... interface and work like torch.utils.data.DataLoader. Right?
• Could you pls give me some advice about the need and the bug? Thanks a lot!

[rusty1s]

I currently cannot run the code, but we have a similar example here which works just fine. As such, I am pretty confident that sampler is fully supported. It might also be related to the usage of drop_last=True.

[rusty1s]

Did this resolve your issues? :)