pgcuts.data

Data loading utilities for edge-pair mini-batch training.

class pgcuts.utils.data.ShuffledRangeDataset(n, k)

Dataset yielding batches of shuffled indices.

Automatically reshuffles when all batches consumed.

Parameters:

• n (int)

• k (int)

__init__(n, k)

Initialize dataset.

Parameters:

• n (int) – Total number of elements.

• k (int) – Batch size (indices per item).

Return type:

None

shuffle()

Reshuffle the permutation.

Return type:

None

Usage:

from torch.utils.data import DataLoader
from pgcuts.data import ShuffledRangeDataset

# n: total number of edges, k: batch size
dataset = ShuffledRangeDataset(n=num_edges, k=8192)
loader = DataLoader(dataset, batch_size=1, num_workers=4)

for batch_idx in loader:
    edge_batch = pairs[batch_idx[0]]