Batched dataloading for contrastive learning

[ziw-liu]

Changes

• Improve throughput of TripletDataModule
  • Use tensorstore for batched reading and decoding of patches
  • Move to a device-agnostic batched transform interface
• Implement transforms that run efficiently on the GPU
  • Batched normalization transforms for both pre-computed global statistics and on-the-fly percentiles scaling
  • Batched affine transform
  • Torch-based RNG for the Gaussian noise transform
• Batched version of ConcatDataModule

Results

Training on a concatenation of 2 data modules from @Soorya19Pradeep’s script takes 6 min/epoch on a small node (16 zen 2 cores, 1 A40 GPU).

Example script:

/hpc/projects/organelle_phenotyping/models/SEC61_TOMM20_G3BP1_Sensor/time_interval/dynaclr_gfp_rfp_Ph/profiling/fit.py

[ziw-liu]

Tests failing due to #277.

[edyoshikun]

There are some transforms that we didn't redefine in Viscy. I added some of them (Decollated and ToDeviced)

I was having trouble running this for multi-channel. (i.e the cases where we want to use two+ channels. Otherwise for 1 channel this works. Solved: this was an issue in my config setting the number of in_channels.

I'll setup a training for the sensor to confirm at least that the is is converging and that we can get similar results.

[edyoshikun]

If we run it as fast_dev_mode=true it runs but we get the error:

Epoch 0: 100%|████████████████████████████████████████████████████████████████████| 1/1 [01:14<00:00,  0.01it/s, loss/train_step=5.200, loss/val_step=4.780, loss/val_epoch=4.780, loss/train_epoch=5.200]`Trainer.fit` stopped: `max_steps=1` reached.                                                                                                                                                             
Epoch 0: 100%|████████████████████████████████████████████████████████████████████| 1/1 [01:14<00:00,  0.01it/s, loss/train_step=5.200, loss/val_step=4.780, loss/val_epoch=4.780, loss/train_epoch=5.200]
Exception ignored in: <function _MultiProcessingDataLoaderIter.__del__ at 0x7fb6592c7920>
Traceback (most recent call last):
  File "/hpc/mydata/eduardo.hirata/.conda/envs/viscy/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1663, in __del__
    self._shutdown_workers()
  File "/hpc/mydata/eduardo.hirata/.conda/envs/viscy/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1651, in _shutdown_workers
    w.terminate()
    ^^^^^^^^^^^

This might be related to the ThreadDataLoader() , I am assuming that we might need to put some guard rails for any additional kwargs that don't apply for threading (i.e not sure if the persistent workers or the prefetch factor are the culprit)

[edyoshikun]

I confirm we are getting a speedboost of at least x2-4 in one node with one A40 gpu. Very active GPU!!! Thanks @ziw-liu

[Soorya19Pradeep]

@edyoshikun , I am unsure if this is realated to your message, but I am getting the following error when trying to run the training with 10 dataloaders.

Sanity Checking: | | 0/? [00:00<?, ?it/s]Traceback (most recent call last):
File "/hpc/projects/organelle_phenotyping/models/SEC61_TOMM20_G3BP1_Sensor/time_interval/dynaclr_gfp_rfp_Ph/fit_concat_dataloader_maxproj_ver2.py", line 349, in
trainer.fit(
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 561, in fit
call._call_and_handle_interrupt(
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 48, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 599, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1012, in _run
results = self._run_stage()
^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1054, in _run_stage
self._run_sanity_check()
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py", line 1083, in _run_sanity_check
val_loop.run()
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/loops/utilities.py", line 179, in _decorator
return loop_run(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 138, in run
batch, batch_idx, dataloader_idx = next(data_fetcher)
^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/loops/fetchers.py", line 134, in next
batch = super().next()
^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/loops/fetchers.py", line 61, in next
batch = next(self.iterator)
^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/utilities/combined_loader.py", line 341, in next
out = next(self._iterator)
^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/lightning/pytorch/utilities/combined_loader.py", line 142, in next
out = next(self.iterators[0])
^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 733, in next
data = self._next_data()
^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1515, in _next_data
return self._process_data(data, worker_id)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1550, in _process_data
data.reraise()
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/_utils.py", line 750, in reraise
raise exception
NotImplementedError: Caught NotImplementedError in DataLoader worker process 0.
Original Traceback (most recent call last):
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/_utils/worker.py", line 349, in _worker_loop
data = fetcher.fetch(index) # type: ignore[possibly-undefined]
^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py", line 52, in
data = [self.dataset[idx] for idx in possibly_batched_index]
~~~~~~~~~~~~^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/dataset.py", line 346, in getitem
return self.datasets[dataset_idx][sample_idx]
~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
File "/hpc/mydata/soorya.pradeep/envs/dynaclr_test/lib/python3.11/site-packages/torch/utils/data/dataset.py", line 59, in getitem
raise NotImplementedError("Subclasses of Dataset should implement getitem.")
NotImplementedError: Subclasses of Dataset should implement getitem.

What would be the solution for this @ziw-liu ? Thanks!

[ziw-liu]

If we run it as fast_dev_mode=true it runs but we get the error:

Epoch 0: 100%|████████████████████████████████████████████████████████████████████| 1/1 [01:14<00:00,  0.01it/s, loss/train_step=5.200, loss/val_step=4.780, loss/val_epoch=4.780, loss/train_epoch=5.200]`Trainer.fit` stopped: `max_steps=1` reached.                                                                                                                                                             
Epoch 0: 100%|████████████████████████████████████████████████████████████████████| 1/1 [01:14<00:00,  0.01it/s, loss/train_step=5.200, loss/val_step=4.780, loss/val_epoch=4.780, loss/train_epoch=5.200]
Exception ignored in: <function _MultiProcessingDataLoaderIter.__del__ at 0x7fb6592c7920>
Traceback (most recent call last):
  File "/hpc/mydata/eduardo.hirata/.conda/envs/viscy/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1663, in __del__
    self._shutdown_workers()
  File "/hpc/mydata/eduardo.hirata/.conda/envs/viscy/lib/python3.11/site-packages/torch/utils/data/dataloader.py", line 1651, in _shutdown_workers
    w.terminate()
    ^^^^^^^^^^^

This might be related to the ThreadDataLoader() , I am assuming that we might need to put some guard rails for any additional kwargs that don't apply for threading (i.e not sure if the persistent workers or the prefetch factor are the culprit)

That's a MONAI bug where they didn't implement the correct teardown method. Doesn't seem to affect our use case though.

[ziw-liu]

I'm seeing some strange behavior. With an independent test dataset (for the encoder) the embedding is similar to what we see in the past:

On the training set, infected cells still cluster, but not in a linearly separated way:

Edit: This also happens with a checkpoint trained before this PR: