fix: download datasets on local rank 0 in multi node (#65)

* fix: download datasets on local rank 0 in multi node

* fix: idist.get_local_rank()

Author: Jeff Yang

templates/gan/datasets.py

@@ -1,5 +1,6 @@
 from torchvision import transforms as T
 from torchvision import datasets as dset
+import ignite.distributed as idist
 
 
 def get_datasets(dataset, dataroot):
@@ -12,6 +13,12 @@ def get_datasets(dataset, dataroot):
     Returns:
         dataset, num_channels
     """
+    local_rank = idist.get_local_rank()
+
+    if local_rank > 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
     resize = T.Resize(64)
     crop = T.CenterCrop(64)
     to_tensor = T.ToTensor()
@@ -42,4 +49,8 @@ def get_datasets(dataset, dataroot):
     else:
         raise RuntimeError(f"Invalid dataset name: {dataset}")
 
+    if local_rank == 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
     return dataset, nc

templates/gan/main.py

@@ -40,16 +40,8 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
     # datasets and dataloaders
     # -----------------------------
 
-    if rank > 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
     train_dataset, num_channels = get_datasets(config.dataset, config.data_path)
 
-    if rank == 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
     train_dataloader = idist.auto_dataloader(
         train_dataset,
         batch_size=config.batch_size,

templates/image_classification/datasets.py

@@ -1,5 +1,6 @@
 from torchvision import datasets
 from torchvision.transforms import Compose, Normalize, Pad, RandomCrop, RandomHorizontalFlip, ToTensor
+import ignite.distributed as idist
 
 train_transform = Compose(
     [
@@ -20,7 +21,17 @@
 
 
 def get_datasets(path):
+    local_rank = idist.get_local_rank()
+
+    if local_rank > 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
     train_ds = datasets.CIFAR10(root=path, train=True, download=True, transform=train_transform)
     eval_ds = datasets.CIFAR10(root=path, train=False, download=True, transform=eval_transform)
 
+    if local_rank == 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
     return train_ds, eval_ds

templates/image_classification/main.py

@@ -36,16 +36,8 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
     # TODO : PLEASE replace `kwargs` with your desirable DataLoader arguments
     # See : https://pytorch.org/ignite/distributed.html#ignite.distributed.auto.auto_dataloader
 
-    if rank > 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
     train_dataset, eval_dataset = get_datasets(path=config.data_path)
 
-    if rank == 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
     train_dataloader = idist.auto_dataloader(
         train_dataset,
         batch_size=config.train_batch_size,
@@ -128,7 +120,9 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
 
     # setup ignite logger only on rank 0
     if rank == 0:
-        logger_handler = get_logger(config=config, train_engine=train_engine, eval_engine=eval_engine, optimizers=optimizer)
+        logger_handler = get_logger(
+            config=config, train_engine=train_engine, eval_engine=eval_engine, optimizers=optimizer
+        )
 
     # -----------------------------------
     # resume from the saved checkpoints

templates/single/datasets.py

@@ -1 +1,21 @@
-# CUSTOM DATASETS AND DATALOADERS GO HERE
+# MAKE SURE YOUR DATASETS ARE DOWNLOADING ON LOCAL_RANK 0.
+
+import ignite.distributed as idist
+
+
+def get_datasets(*args, **kwargs):
+    local_rank = idist.get_local_rank()
+
+    if local_rank > 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
+    # CUSTOM DATASETS GO HERE
+    train_dataset = ...
+    eval_dataset = ...
+
+    if local_rank == 0:
+        # Ensure that only rank 0 download the dataset
+        idist.barrier()
+
+    return train_dataset, eval_dataset

templates/single/main.py

@@ -11,6 +11,7 @@
 from ignite.engine.events import Events
 from ignite.utils import manual_seed
 
+from datasets import get_datasets
 from trainers import create_trainers, TrainEvents
 from handlers import get_handlers, get_logger
 from utils import setup_logging, log_metrics, log_basic_info, initialize, resume_from
@@ -34,16 +35,7 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
     # TODO : PLEASE replace `kwargs` with your desirable DataLoader arguments
     # See : https://pytorch.org/ignite/distributed.html#ignite.distributed.auto.auto_dataloader
 
-    if rank > 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
-
-    train_dataset = ...
-    eval_dataset = ...
-
-    if rank == 0:
-        # Ensure that only rank 0 download the dataset
-        idist.barrier()
+    train_dataset, eval_dataset = get_datasets()
 
     train_dataloader = idist.auto_dataloader(train_dataset, **kwargs)
     eval_dataloader = idist.auto_dataloader(eval_dataset, **kwargs)
@@ -104,7 +96,9 @@ def run(local_rank: int, config: Any, *args: Any, **kwargs: Any):
 
     # setup ignite logger only on rank 0
     if rank == 0:
-        logger_handler = get_logger(config=config, train_engine=train_engine, eval_engine=eval_engine, optimizers=optimizer)
+        logger_handler = get_logger(
+            config=config, train_engine=train_engine, eval_engine=eval_engine, optimizers=optimizer
+        )
 
     # -----------------------------------
     # resume from the saved checkpoints