Update distributed training for MetaTensor (Project-MONAI#777)

* [DLMED] update brats ddp

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] restore dist for compatible

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] update ddp examples

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] update ignite workflows

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] update smart cache ddp

Signed-off-by: Nic Ma <[email protected]>

* [DLMED] update horovod ddp

Signed-off-by: Nic Ma <[email protected]>

Author: Nic-Ma

Diff for: acceleration/distributed_training/brats_training_ddp.py

@@ -57,7 +57,6 @@
 """
 
 import argparse
-import numpy as np
 import os
 import sys
 import time
@@ -89,8 +88,6 @@
     RandSpatialCropd,
     Spacingd,
     ToDeviced,
-    EnsureTyped,
-    EnsureType,
 )
 from monai.utils import set_determinism
 
@@ -110,16 +107,16 @@ def __call__(self, data):
         for key in self.keys:
             result = []
             # merge label 2 and label 3 to construct TC
-            result.append(np.logical_or(d[key] == 2, d[key] == 3))
+            result.append(torch.logical_or(d[key] == 2, d[key] == 3))
             # merge labels 1, 2 and 3 to construct WT
             result.append(
-                np.logical_or(
-                    np.logical_or(d[key] == 2, d[key] == 3), d[key] == 1
+                torch.logical_or(
+                    torch.logical_or(d[key] == 2, d[key] == 3), d[key] == 1
                 )
             )
             # label 2 is ET
             result.append(d[key] == 2)
-            d[key] = np.stack(result, axis=0).astype(np.float32)
+            d[key] = torch.stack(result, dim=0)
         return d
 
 
@@ -132,7 +129,7 @@ def __init__(
         self,
         root_dir,
         section,
-        transform=LoadImaged(["image", "label"]),
+        transform=None,
         cache_rate=1.0,
         num_workers=0,
         shuffle=False,
@@ -187,6 +184,7 @@ def main_worker(args):
         [
             # load 4 Nifti images and stack them together
             LoadImaged(keys=["image", "label"]),
+            ToDeviced(keys=["image", "label"], device=device),
             EnsureChannelFirstd(keys="image"),
             ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
             Orientationd(keys=["image", "label"], axcodes="RAS"),
@@ -195,8 +193,6 @@ def main_worker(args):
                 pixdim=(1.0, 1.0, 1.0),
                 mode=("bilinear", "nearest"),
             ),
-            EnsureTyped(keys=["image", "label"]),
-            ToDeviced(keys=["image", "label"], device=device),
             RandSpatialCropd(keys=["image", "label"], roi_size=[224, 224, 144], random_size=False),
             RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
             RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
@@ -223,6 +219,7 @@ def main_worker(args):
     val_transforms = Compose(
         [
             LoadImaged(keys=["image", "label"]),
+            ToDeviced(keys=["image", "label"], device=device),
             EnsureChannelFirstd(keys="image"),
             ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
             Orientationd(keys=["image", "label"], axcodes="RAS"),
@@ -232,8 +229,6 @@ def main_worker(args):
                 mode=("bilinear", "nearest"),
             ),
             NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
-            EnsureTyped(keys=["image", "label"]),
-            ToDeviced(keys=["image", "label"], device=device),
         ]
     )
     val_ds = BratsCacheDataset(
@@ -283,7 +278,7 @@ def main_worker(args):
     dice_metric = DiceMetric(include_background=True, reduction="mean")
     dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")
 
-    post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
 
     # start a typical PyTorch training
     best_metric = -1

Diff for: acceleration/distributed_training/unet_evaluation_ddp.py

@@ -62,7 +62,7 @@
 from monai.data import DataLoader, Dataset, create_test_image_3d, DistributedSampler, decollate_batch
 from monai.inferers import sliding_window_inference
 from monai.metrics import DiceMetric
-from monai.transforms import Activations, AsChannelFirstd, AsDiscrete, Compose, LoadImaged, ScaleIntensityd, EnsureTyped, EnsureType
+from monai.transforms import Activations, AsChannelFirstd, AsDiscrete, Compose, LoadImaged, ScaleIntensityd
 
 
 def evaluate(args):
@@ -92,7 +92,6 @@ def evaluate(args):
             LoadImaged(keys=["img", "seg"]),
             AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
             ScaleIntensityd(keys="img"),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
 
@@ -103,7 +102,7 @@ def evaluate(args):
     # sliding window inference need to input 1 image in every iteration
     val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=2, pin_memory=True, sampler=val_sampler)
     dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
-    post_trans = Compose([EnsureType(), Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
+    post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
     # create UNet, DiceLoss and Adam optimizer
     device = torch.device(f"cuda:{args.local_rank}")
     torch.cuda.set_device(device)

Diff for: acceleration/distributed_training/unet_evaluation_horovod.py

@@ -35,7 +35,7 @@
     Example script to execute this program, only need to run on the master node:
     `horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python unet_evaluation_horovod.py -d "./testdata"`
 
-    This example was tested with [Ubuntu 16.04/20.04], [NCCL 2.6.3], [horovod 0.19.5].
+    This example was tested with [Ubuntu 16.04/20.04], [NCCL 2.6.3], [horovod 0.25.0].
 
 Referring to: https://github.com/horovod/horovod/blob/master/examples/pytorch_mnist.py
 
@@ -56,7 +56,7 @@
 from monai.data import DataLoader, Dataset, create_test_image_3d, decollate_batch
 from monai.inferers import sliding_window_inference
 from monai.metrics import DiceMetric
-from monai.transforms import Activations, AsChannelFirstd, AsDiscrete, Compose, LoadImaged, ScaleIntensityd, EnsureTyped, EnsureType
+from monai.transforms import Activations, AsChannelFirstd, AsDiscrete, Compose, LoadImaged, ScaleIntensityd, EnsureType
 
 
 def evaluate(args):
@@ -88,7 +88,6 @@ def evaluate(args):
             LoadImaged(keys=["img", "seg"]),
             AsChannelFirstd(keys=["img", "seg"], channel_dim=-1),
             ScaleIntensityd(keys="img"),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
 
@@ -156,7 +155,7 @@ def main():
     evaluate(args=args)
 
 
-# Example script to execute this program only on the master node:
+# Example script to execute this program on 4 nodes (only need to run below command on the master node):
 # horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python unet_evaluation_horovod.py -d "./testdata"
 if __name__ == "__main__":
     main()

Diff for: acceleration/distributed_training/unet_evaluation_workflows.py

@@ -79,7 +79,6 @@
     KeepLargestConnectedComponentd,
     LoadImaged,
     ScaleIntensityd,
-    EnsureTyped,
     SaveImaged,
 )
 
@@ -113,7 +112,6 @@ def evaluate(args):
             LoadImaged(keys=["image", "label"]),
             AsChannelFirstd(keys=["image", "label"], channel_dim=-1),
             ScaleIntensityd(keys="image"),
-            EnsureTyped(keys=["image", "label"]),
         ]
     )
 

Diff for: acceleration/distributed_training/unet_training_ddp.py

@@ -67,7 +67,6 @@
     RandCropByPosNegLabeld,
     RandRotate90d,
     ScaleIntensityd,
-    EnsureTyped,
 )
 
 
@@ -106,7 +105,6 @@ def train(args):
                 keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
             ),
             RandRotate90d(keys=["img", "seg"], prob=0.5, spatial_axes=[0, 2]),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
 

Diff for: acceleration/distributed_training/unet_training_horovod.py

@@ -39,7 +39,7 @@
     Example script to execute this program, only need to run on the master node:
     `horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python unet_training_horovod.py -d "./testdata"`
 
-    This example was tested with [Ubuntu 16.04/20.04], [NCCL 2.6.3], [horovod 0.19.5].
+    This example was tested with [Ubuntu 16.04/20.04], [NCCL 2.6.3], [horovod 0.25.0].
 
 Referring to: https://github.com/horovod/horovod/blob/master/examples/pytorch_mnist.py
 
@@ -66,7 +66,6 @@
     RandCropByPosNegLabeld,
     RandRotate90d,
     ScaleIntensityd,
-    EnsureTyped,
 )
 
 
@@ -106,7 +105,6 @@ def train(args):
                 keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
             ),
             RandRotate90d(keys=["img", "seg"], prob=0.5, spatial_axes=[0, 2]),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
 
@@ -188,7 +186,7 @@ def main():
     train(args=args)
 
 
-# Example script to execute this program only on the master node:
+# Example script to execute this program on 4 nodes (only need to run below command on the master node):
 # horovodrun -np 16 -H server1:4,server2:4,server3:4,server4:4 python unet_training_horovod.py -d "./testdata"
 if __name__ == "__main__":
     main()

Diff for: acceleration/distributed_training/unet_training_smartcache.py

@@ -73,7 +73,6 @@
     RandCropByPosNegLabeld,
     RandRotate90d,
     ScaleIntensityd,
-    EnsureTyped,
 )
 
 
@@ -112,7 +111,6 @@ def train(args):
                 keys=["img", "seg"], label_key="seg", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
             ),
             RandRotate90d(keys=["img", "seg"], prob=0.5, spatial_axes=[0, 2]),
-            EnsureTyped(keys=["img", "seg"]),
         ]
     )
 

Diff for: acceleration/distributed_training/unet_training_workflows.py

@@ -81,7 +81,6 @@
     RandCropByPosNegLabeld,
     RandRotate90d,
     ScaleIntensityd,
-    EnsureTyped,
 )
 
 
@@ -118,7 +117,6 @@ def train(args):
                 keys=["image", "label"], label_key="label", spatial_size=[96, 96, 96], pos=1, neg=1, num_samples=4
             ),
             RandRotate90d(keys=["image", "label"], prob=0.5, spatial_axes=[0, 2]),
-            EnsureTyped(keys=["image", "label"]),
         ]
     )
 
@@ -155,7 +153,6 @@ def train(args):
 
     train_post_transforms = Compose(
         [
-            EnsureTyped(keys="pred"),
             Activationsd(keys="pred", sigmoid=True),
             AsDiscreted(keys="pred", threshold=0.5),
             KeepLargestConnectedComponentd(keys="pred", applied_labels=[1]),
@@ -198,7 +195,6 @@ def main():
     train(args=args)
 
 
-
 # python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_PER_NODE
 #        --nnodes=NUM_NODES --node_rank=INDEX_CURRENT_NODE
 #        --master_addr="192.168.1.1" --master_port=1234