MolecularDiffusion.data.dataloader¶

Classes¶

`DataLoader`	Extended data loader for batching graph structured data.
`DataQueue`

Functions¶

`estimate_batch_size`(nmax[, vram_size])
`graph_collate`(batch)	Convert any list of same nested container into a container of tensors.
`pointcloud_collate`([vram_size])
`pointcloud_collate_v0`(batch[, to_keep])	Convert any list of same nested container into a container of tensors.

Module Contents¶

class MolecularDiffusion.data.dataloader.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0, collate_fn=graph_collate, **kwargs)¶

Bases: torch.utils.data.DataLoader

Extended data loader for batching graph structured data.

See torch.utils.data.DataLoader for more details.

Parameters:

dataset (Dataset) – dataset from which to load the data
batch_size (int, optional) – how many samples per batch to load
shuffle (bool, optional) – set to True to have the data reshuffled at every epoch
sampler (Sampler, optional) – sampler that draws single sample from the dataset
batch_sampler (Sampler, optional) – sampler that draws a mini-batch of data from the dataset
num_workers (int, optional) – how many subprocesses to use for data loading
collate_fn (callable, optional) – merge a list of samples into a mini-batch
kwargs – keyword arguments for torch.utils.data.DataLoader

class MolecularDiffusion.data.dataloader.DataQueue¶

Bases: torch.utils.data.Dataset

append(item)¶

pop()¶

queue¶

MolecularDiffusion.data.dataloader.estimate_batch_size(nmax, vram_size=40)¶

MolecularDiffusion.data.dataloader.graph_collate(batch)¶

Convert any list of same nested container into a container of tensors.

Parameters:: batch (list) – list of samples with the same nested container

MolecularDiffusion.data.dataloader.pointcloud_collate(vram_size=40)¶

MolecularDiffusion.data.dataloader.pointcloud_collate_v0(batch, to_keep=None)¶

Convert any list of same nested container into a container of tensors.

Parameters:: batch (list) – list of samples with the same nested container