Preload small datasets to memory in Custom dataset

[RaulPPelaez]

While replicating results from the torchmd-protein-thermodynamics repository, I experienced sluggish training speeds and low GPU usage (meaning sitting at 0% and briefly going to 100% at each iteration) using the following configuration file:

activation: tanh
batch_size: 1024
inference_batch_size: 1024
dataset: Custom
coord_files: "chignolin_ca_coords.npy"
embed_files: "chignolin_ca_embeddings.npy"
force_files: "chignolin_ca_deltaforces.npy"
cutoff_upper: 12.0
cutoff_lower: 3.0
derivative: true
early_stopping_patience: 30
embedding_dimension: 128
lr: 0.0005
lr_factor: 0.8
lr_min: 1.0e-06
lr_patience: 10
lr_warmup_steps: 0
model: graph-network
neighbor_embedding: false
ngpus: -1
num_epochs: 100
num_layers: 4
num_nodes: 1
num_rbf: 18
num_workers: 4
rbf_type: expnorm
save_interval: 2
seed: 1
test_interval: 2
test_size: 0.1
trainable_rbf: true
val_size: 0.05
weight_decay: 0.0

The referenced files take approx 300mb.
Playing around with num_workers and batch size did not help.
Upon investigation, the issue arose from the Custom dataset's I/O-bound get method, which reads from disk every time it is invoked, causing the low GPU usage and slow training.

To resolve this, I implemented a preloading feature that loads the complete dataset into system memory if its size is below a user-configurable threshold, set by default at 1GB. The data is stored as PyTorch tensors, facilitating compatibility with multi-threaded data loaders (num_workers). Notably, this approach does not inflate RAM usage when increasing the number of workers.

This optimization led to a x20 speedup in training time for this specific setup.

OTOH I tweaked the DataLoader options a bit.

[RaulPPelaez]

@AntonioMirarchi please review!

[peastman]

The Custom dataset class is incredibly inefficient. It reloads the whole file every time get() is called to retrieve a single sample. Some intelligent caching would help. But a much better choice is to use the HDF5 dataset class. It is far more efficient.

[peastman]

In fact, possibly we should just make Custom create a temporary HDF5 file on startup and then load from it.

[RaulPPelaez]

I really like this idea Peter, thanks!
I do not like Custom at all either, but it offers some convenience/simplicity that makes people choose it so I think it is worth improving.
I implemented so that the user can instruct Custom to use HDF5 under the hood, lets see how it goes.

[giadefa]

h5py are slower than mmap arrays. Maybe I am missing the narrative but why are we doing something different from what we are already doing in other dataloaders? g

…

On Mon, Oct 16, 2023 at 11:05 AM Antonio Mirarchi ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In torchmdnet/datasets/custom.py <#235 (comment)>: > + with h5py.File(hdf5_dataset, "w") as f: + for i in range(len(files["pos"])): + # Create a group for each file + coord_data = np.load(files["pos"][i]) + embed_data = np.load(files["z"][i]).astype(int) + group = f.create_group(str(i)) + num_samples = coord_data.shape[0] + group["pos"] = coord_data + group["types"] = np.tile(embed_data, (num_samples, 1)) + if "y" in files: + energy_data = np.load(files["y"][i]) + group["energy"] = energy_data + if "neg_dy" in files: + force_data = np.load(files["neg_dy"][i]) + group["forces"] = force_data It's just the "proper way", in theory you should be able to move across path in the hdf5 file (it's not this case) . For example if you use "create_dataset" you can retrieve the pos dataset using: f = h5py.File(hdf5_dataset, "r")f["0/pos"] While if you use the "dictionary way" you get this error: KeyError: "Unable to open object (object 'pos' doesn't exist)" — Reply to this email directly, view it on GitHub <#235 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AB3KUOUKJXJDPK7SBVMBLRTX7T2FNANCNFSM6AAAAAA54GA4WI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[RaulPPelaez]

This PR exists because Custom was calling np.load() at each call of get(). Even with mmap mode this was really slowing down trainings (it is I/O bound in the end...).
I changed it so that:
1- If the dataset is small enough just load it all into RAM
2- Otherwise store the references to the mmap arrays instead of reloading the file each time.
Alternatively, I added an option to transform the Custom files to HDF5, which seems to be just a little bit slower than mmap.
I went ahead and also implemented the same load to RAM functionality in HDF5.

[RaulPPelaez]

This is ready. @stefdoerr could you please review?

[stefdoerr]

I think @raimis should review this since he worked on mmaps before. I am not too qualified, so I just commented on style and minor bugs

[RaulPPelaez]

I would like to merge this one before next release, @raimis could you take a look? Thanks

[guillemsimeon]

is this ready to be shipped?

[RaulPPelaez]

Yes!

[guillemsimeon]

perf

[guillemsimeon]

we should merge these