torchcell.loader.dense_padding_data_loader.DensePaddingDataLoader

class DensePaddingDataLoader(dataset: Dataset | Sequence[BaseData] | DatasetAdapter, batch_size: int = 1, shuffle: bool = False, follow_batch: list[str] | None = None, exclude_keys: list[str] | None = None, **kwargs: Any)[source]

Bases: DataLoader[Any]

Data loader that merges dataset objects into mini-batches via dense padding.

Merges data objects from a torch_geometric.data.Dataset into a mini-batch using padding along a new dimension. Data objects can be either of type ~torch_geometric.data.Data or ~torch_geometric.data.HeteroData.

Parameters:

  • dataset (Dataset) – The dataset from which to load the data.

  • batch_size (int, optional) – How many samples per batch to load. (default: 1)

  • shuffle (bool, optional) – If set to True, the data will be reshuffled at every epoch. (default: False)

  • follow_batch (List[str], optional) – Creates assignment batch vectors for each key in the list. (default: None)

  • exclude_keys (List[str], optional) – Will exclude each key in the list. (default: None)

  • **kwargs (optional) – Additional arguments of torch.utils.data.DataLoader.