data.AtomsDataModule

class data.AtomsDataModule(*args: Any, **kwargs: Any)

A general LightningDataModule for SchNetPack datasets.

Parameters:

• datapath – path to dataset

• batch_size – (train) batch size

• num_train – number of training examples (absolute or relative)

• num_val – number of validation examples (absolute or relative)

• num_test – number of test examples (absolute or relative)

• split_file – path to npz file with data partitions

• format – dataset format

• load_properties – subset of properties to load

• val_batch_size – validation batch size. If None, use test_batch_size, then batch_size.

• test_batch_size – test batch size. If None, use val_batch_size, then batch_size.

• transforms – Preprocessing transform applied to each system separately before batching.

• train_transforms – Overrides transform_fn for training.

• val_transforms – Overrides transform_fn for validation.

• test_transforms – Overrides transform_fn for testing.

• train_sampler_cls – type of torch training sampler. This is by default wrapped into a torch.utils.data.BatchSampler.

• train_sampler_args – dict of train_sampler keyword arguments.

• num_workers – Number of data loader workers.

• num_val_workers – Number of validation data loader workers (overrides num_workers).

• num_test_workers – Number of test data loader workers (overrides num_workers).

• property_units – Dictionary from property to corresponding unit as a string (eV, kcal/mol, …).

• distance_unit – Unit of the atom positions and cell as a string (Ang, Bohr, …).

• data_workdir – Copy data here as part of setup, e.g. to a local file system for faster performance.

• cleanup_workdir_stage – Determines after which stage to remove the data workdir

• splitting – Method to generate train/validation/test partitions (default: RandomSplit)

• pin_memory – If true, pin memory of loaded data to GPU. Default: Will be set to true, when GPUs are used.