Sharded distributed sampler for cached dataloading in DDP

[ziw-liu]

Add a distributed sampler that only permutes index within ranks, improving cache hit rate in DDP.

See viscy/scripts/shared_dict.py for usage.

Also includes changes from #196.

[ziw-liu]

Example output:

GPU available: True (cuda), used: False TPU available: False, using: 0 TPU cores HPU available: False, using: 0 HPUs /hpc/mydata/ziwen.liu/anaconda/2022.05/x86_64/envs/viscy/lib/python3.11/site-packages/lightning/pytorch/trainer/setup.py:177: GPU available but not used. You can set it by doing `Trainer(accelerator='gpu')`. Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/3 Initializing distributed: GLOBAL_RANK: 2, MEMBER: 3/3 Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/3 ---------------------------------------------------------------------------------------------------- distributed_backend=gloo All distributed processes registered. Starting with 3 processes ----------------------------------------------------------------------------------------------------

=== Initializing cache pool for rank 0 ===
=== Initializing cache pool for rank 1 ===
=== Initializing cache pool for rank 2 ===

| Name | Type | Params | Mode

0 | layer | Linear | 2 | train

2 Trainable params
0 Non-trainable params
2 Total params
0.000 Total estimated model params size (MB)
1 Modules in train mode
0 Modules in eval mode

• Adding 31 to cache dict on rank 1
• Adding 32 to cache dict on rank 2
• Adding 38 to cache dict on rank 2
• Adding 42 to cache dict on rank 0
• Adding 30 to cache dict on rank 0
• Adding 36 to cache dict on rank 0
• Adding 37 to cache dict on rank 1
• Adding 43 to cache dict on rank 1
• Adding 48 to cache dict on rank 0
• Adding 34 to cache dict on rank 1
• Adding 49 to cache dict on rank 1
• Adding 30 to cache dict on rank 2
• Adding 44 to cache dict on rank 2
• Adding 41 to cache dict on rank 2
• Adding 35 to cache dict on rank 2
• Adding 40 to cache dict on rank 1
• Adding 46 to cache dict on rank 1
• Adding 39 to cache dict on rank 0
• Adding 33 to cache dict on rank 0
• Adding 47 to cache dict on rank 2
• Adding 45 to cache dict on rank 0
• Adding 24 to cache dict on rank 2
• Adding 13 to cache dict on rank 1
• Adding 0 to cache dict on rank 0
• Adding 20 to cache dict on rank 2
• Adding 4 to cache dict on rank 0
• Adding 29 to cache dict on rank 2
• Adding 19 to cache dict on rank 1
• Adding 26 to cache dict on rank 2
• Adding 28 to cache dict on rank 2
  === Starting training ===
  === Starting training epoch 0 ===
• Adding 8 to cache dict on rank 0
• Adding 15 to cache dict on rank 1
• Adding 3 to cache dict on rank 0
• Adding 21 to cache dict on rank 2
• Adding 11 to cache dict on rank 1
• Adding 7 to cache dict on rank 0
• Adding 23 to cache dict on rank 2
• Adding 27 to cache dict on rank 2
• Adding 22 to cache dict on rank 2
• Adding 1 to cache dict on rank 0
• Adding 9 to cache dict on rank 0
• Adding 5 to cache dict on rank 0
• Adding 17 to cache dict on rank 1
• Adding 6 to cache dict on rank 0
• Adding 18 to cache dict on rank 1
• Adding 16 to cache dict on rank 1
• Adding 14 to cache dict on rank 1
• Adding 10 to cache dict on rank 1
• Adding 25 to cache dict on rank 2
• Adding 2 to cache dict on rank 0
• Adding 12 to cache dict on rank 1
  === Starting training epoch 1 ===
  === Starting training epoch 2 ===
  === Starting training epoch 3 ===
  === Starting training epoch 4 ===
  Trainer.fit stopped: max_epochs=5 reached.

[ziw-liu]

The FcmaeUNet class can now do both pre-training and fine-tuning. This is controlled via the pretraining flag in model_config. This allows keeping the old VSUNet with its integration with the sliding window dataset, but also enable GPU-accelerated augmentations for virtual staining training.

[ziw-liu]

@edyoshikun do you want to keep improving the prototype in hcs_ram.py or we can remove it?

[edyoshikun]

remove the hcs_ram.py

[edyoshikun]

This LGTM. The pending thing would be to properly write the augmentations now. Happy to merge this first and then add the augmentations.

[ziw-liu]

properly write the augmentations now

Can you elaborate?

[edyoshikun]

I mean the tiling in 3D. I was thinking of the neuromast dataset or mantis datasets where if the XY dimension is less than 1 patch, then we only crop/tile the top left corners. The equivalent thing would happen in the Z dimension.

[ziw-liu]

I mean the tiling in 3D. I was thinking of the neuromast dataset or mantis datasets where if the XY dimension is less than 1 patch, then we only crop/tile the top left corners. The equivalent thing would happen in the Z dimension.

IIRC this only affects validation? Let's open an issue and fix the tiling transform separately.

[edyoshikun]

I have super minor comments. Otherwise I've tested this with the phase pertaining and works well.

Let's fix the tiling in a separate issue+PR.

[edyoshikun]

Minor preference but when I see ratios I typically think of values 0-1. Like the masking ratio, the train/val ratio, etc.

[ziw-liu]

What would be a better name?

[ziw-liu]

Let's keep this for now as I am the only user at the moment.