Spaces:
Paused
Paused
File size: 2,329 Bytes
5716801 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import warnings
from torch.utils.data import DataLoader
from chemprop.data.collate import collate_batch, collate_multicomponent
from chemprop.data.datasets import MoleculeDataset, MulticomponentDataset, ReactionDataset
from chemprop.data.samplers import ClassBalanceSampler, SeededSampler
def build_dataloader(
dataset: MoleculeDataset | ReactionDataset | MulticomponentDataset,
batch_size: int = 64,
num_workers: int = 0,
class_balance: bool = False,
seed: int | None = None,
shuffle: bool = True,
**kwargs,
):
"""Return a :obj:`~torch.utils.data.DataLoader` for :class:`MolGraphDataset`\s
Parameters
----------
dataset : MoleculeDataset | ReactionDataset | MulticomponentDataset
The dataset containing the molecules or reactions to load.
batch_size : int, default=64
the batch size to load.
num_workers : int, default=0
the number of workers used to build batches.
class_balance : bool, default=False
Whether to perform class balancing (i.e., use an equal number of positive and negative
molecules). Class balance is only available for single task classification datasets. Set
shuffle to True in order to get a random subset of the larger class.
seed : int, default=None
the random seed to use for shuffling (only used when `shuffle` is `True`).
shuffle : bool, default=False
whether to shuffle the data during sampling.
"""
if class_balance:
sampler = ClassBalanceSampler(dataset.Y, seed, shuffle)
elif shuffle and seed is not None:
sampler = SeededSampler(len(dataset), seed)
else:
sampler = None
if isinstance(dataset, MulticomponentDataset):
collate_fn = collate_multicomponent
else:
collate_fn = collate_batch
if len(dataset) % batch_size == 1:
warnings.warn(
f"Dropping last batch of size 1 to avoid issues with batch normalization \
(dataset size = {len(dataset)}, batch_size = {batch_size})"
)
drop_last = True
else:
drop_last = False
return DataLoader(
dataset,
batch_size,
sampler is None and shuffle,
sampler,
num_workers=num_workers,
collate_fn=collate_fn,
drop_last=drop_last,
**kwargs,
)
|