Batching in NeighborSampler

[mailforharry98]

Hello
I want to create balanced batches for training but could not figure out where is batching done in the NeighborSampler code.

[wsad1]

Are you referring to torch_geometric.nn.loader.NeighborSampler? Could you elaborate on what you mean by balanced batches?

[mailforharry98]

Yes, and by balanced batches I mean a batch containing an equal number of samples from every class.

[wsad1]

Firstly use NeighborLoader, since NeighborSampler is depreciated.
Also check out ImbalancedSampler this works with NeighborLoader and makes sure the classes are balanced in a batch.

[SulafahKhaled]

If I were to use a stratifiedKfold sampler to sample before using NeighborLoader, wouldn't that solve the problem with the imbalance as well as provide more scientifically accurate results?

[mzamini92]

try something like this:

import torch
from torch_geometric.datasets import Planetoid
from torch_geometric.loader import NeighborSampler, WeightedRandomSampler

dataset = Planetoid(root='.', name='Cora')

# Get the class labels for the dataset.
labels = dataset.y

# Calculate the class weights.
class_weights = torch.ones(len(labels))
class_weights[labels == 0] = 2

# Create a WeightedRandomSampler instance.
sampler = WeightedRandomSampler(dataset, class_weights)

# Iterate over the batches.
for batch in sampler:
    print(batch)

This code will create a WeightedRandomSampler instance and then iterate over the batches. Each batch will contain an equal number of samples from the two classes.