Repository: Project-MONAI/MONAI

Issue:
Support more backends in distributed operations
The current version of MONAI utilizes `torch.distributed` for `all_gather`, `all_reduce` etc., e.g. in **line 409** below. However, this will raise an Error when I launch the parallel programs in `ignite.distributed.Parallel` context with `horovod` or `xla` backend: `RuntimeError: Default process group has not been initialized, please make sure to call init_process_group.` Is it better to replace all the `torch.distributed` operations with `ignite.distributed` ones? They have naive support for all the backends.

https://github.com/Project-MONAI/MONAI/blob/1a707f0226bfe10c000cc54ac87b35d50973e221/monai/utils/misc.py#L359-L423


Discussion:

Patch:
diff --git a/monai/config/deviceconfig.py b/monai/config/deviceconfig.py
--- a/monai/config/deviceconfig.py
+++ b/monai/config/deviceconfig.py
@@ -19,7 +19,7 @@
 import torch
 
 import monai
-from monai.utils import OptionalImportError, get_package_version, optional_import
+from monai.utils.module import OptionalImportError, get_package_version, optional_import
 
 try:
     import itk  # type: ignore
diff --git a/monai/handlers/checkpoint_saver.py b/monai/handlers/checkpoint_saver.py
--- a/monai/handlers/checkpoint_saver.py
+++ b/monai/handlers/checkpoint_saver.py
@@ -11,7 +11,7 @@
 
 import logging
 import warnings
-from typing import TYPE_CHECKING, Dict, Optional, Union
+from typing import TYPE_CHECKING, Dict, Optional
 
 from monai.config import IgniteInfo
 from monai.utils import min_version, optional_import
@@ -62,9 +62,9 @@ class CheckpointSaver:
             typically, it's used to resume training and compare current metric with previous N values.
         key_metric_greater_or_equal: if `True`, the latest equally scored model is stored. Otherwise,
             save the the first equally scored model. default to `False`.
-        key_metric_score_sign: sign of the score, available value: `[1.0, -1.0, 1, -1]`. For error-like metrics,
-            e.g. smaller is better, a negative score sign should be used (objects with larger score are retained).
-            default to `1.0`.
+        key_metric_negative_sign: whether adding a negative sign to the metric score to compare metrics,
+            because for error-like metrics, smaller is better(objects with larger score are retained).
+            default to `False`.
         epoch_level: save checkpoint during training for every N epochs or every N iterations.
             `True` is epoch level, `False` is iteration level.
         save_interval: save checkpoint every N epochs, default is 0 to save no checkpoint.
@@ -96,7 +96,7 @@ def __init__(
         key_metric_filename: Optional[str] = None,
         key_metric_save_state: bool = False,
         key_metric_greater_or_equal: bool = False,
-        key_metric_score_sign: Union[float, int] = 1.0,
+        key_metric_negative_sign: bool = False,
         epoch_level: bool = True,
         save_interval: int = 0,
         n_saved: Optional[int] = None,
@@ -159,10 +159,8 @@ def _score_func(engine: Engine):
                     raise ValueError(
                         f"Incompatible values: save_key_metric=True and key_metric_name={key_metric_name}."
                     )
-                if key_metric_score_sign not in (1.0, -1.0, 1, -1):
-                    raise ValueError("available value of key_metric_score_sign: `[1.0, -1.0, 1, -1]`.")
 
-                return key_metric_score_sign * engine.state.metrics[metric_name]
+                return (-1 if key_metric_negative_sign else 1) * engine.state.metrics[metric_name]
 
             if key_metric_filename is not None and key_metric_n_saved > 1:
                 raise ValueError("if using fixed filename to save the best metric model, we should only save 1 model.")
diff --git a/monai/utils/__init__.py b/monai/utils/__init__.py
--- a/monai/utils/__init__.py
+++ b/monai/utils/__init__.py
@@ -13,6 +13,7 @@
 from .aliases import alias, resolve_name
 from .decorators import MethodReplacer, RestartGenerator
 from .deprecated import DeprecatedError, deprecated, deprecated_arg
+from .dist import evenly_divisible_all_gather, get_dist_device, string_list_all_gather
 from .enums import (
     Activation,
     Average,
@@ -44,10 +45,8 @@
     ensure_tuple,
     ensure_tuple_rep,
     ensure_tuple_size,
-    evenly_divisible_all_gather,
     fall_back_tuple,
     first,
-    get_dist_device,
     get_seed,
     is_scalar,
     is_scalar_tensor,
@@ -56,7 +55,6 @@
     progress_bar,
     set_determinism,
     star_zip_with,
-    string_list_all_gather,
     zip_with,
 )
 from .module import (
diff --git a/monai/utils/dist.py b/monai/utils/dist.py
new file mode 100644
--- /dev/null
+++ b/monai/utils/dist.py
@@ -0,0 +1,154 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+import torch
+import torch.distributed as dist
+
+from monai.config import IgniteInfo
+from monai.utils.module import min_version, optional_import
+
+idist, has_ignite = optional_import("ignite", IgniteInfo.OPT_IMPORT_VERSION, min_version, "distributed")
+
+__all__ = ["get_dist_device", "evenly_divisible_all_gather", "string_list_all_gather"]
+
+
+def get_dist_device():
+    """
+    Get the expected target device in the native PyTorch distributed data parallel.
+    For NCCL backend, return GPU device of current process.
+    For GLOO backend, return CPU.
+    For any other backends, return None as the default, tensor.to(None) will not change the device.
+
+    """
+    if dist.is_initialized():
+        backend = dist.get_backend()
+        if backend == "nccl" and torch.cuda.is_available():
+            return torch.device(f"cuda:{torch.cuda.current_device()}")
+        elif backend == "gloo":
+            return torch.device("cpu")
+    return None
+
+
+def evenly_divisible_all_gather(data: torch.Tensor, concat: bool = True):
+    """
+    Utility function for distributed data parallel to pad at first dim to make it evenly divisible and all_gather.
+    The input data of every rank should have the same number of dimensions, only the first dim can be different.
+
+    Note: If has ignite installed, will execute based on ignite distributed APIs, otherwise, if the native
+    PyTorch distributed group initialized, will execute based on native PyTorch distributed APIs.
+
+    Args:
+        data: source tensor to pad and execute all_gather in distributed data parallel.
+        concat: whether to concat the gathered list to be a Tensor, if False, return a list
+            of Tensors, similar behavior as torch.distributed.all_gather(). default to True.
+
+    Note:
+        The input data on different ranks must have exactly same `dtype`.
+
+    """
+    if not isinstance(data, torch.Tensor):
+        raise ValueError("input data must be PyTorch Tensor.")
+    # data of all the ranks must have same number of dimensions
+    ndims = data.ndimension()
+    length: int = data.shape[0] if ndims > 0 else 1
+
+    def _torch_all_gather(data: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Implementation based on native PyTorch distributed data parallel APIs.
+
+        """
+        device = get_dist_device()
+        orig_device = data.device
+        data = data.to(device)
+        data = data.unsqueeze(0) if ndims == 0 else data
+
+        # make sure the data is evenly-divisible on multi-GPUs
+        length_tensor = torch.as_tensor([length], device=device)
+        all_lens = [torch.zeros_like(length_tensor) for _ in range(dist.get_world_size())]
+        dist.all_gather(all_lens, length_tensor)
+        all_lens_: List[int] = [int(i.item()) for i in all_lens]
+
+        max_len: int = max(all_lens_)
+        if length < max_len:
+            size = [max_len - length] + list(data.shape[1:])
+            data = torch.cat([data, data.new_full(size, 0)], dim=0)
+        # all gather across all processes
+        output = [torch.zeros_like(data) for _ in range(dist.get_world_size())]
+        dist.all_gather(output, data)
+        # remove the padding items, if all the input data doesn't have batch dim, suqeeze the first dim
+        return [(o.squeeze(0) if ndims == 0 else o[:l, ...]).to(orig_device) for o, l in zip(output, all_lens_)]
+
+    def _ignite_all_gather(data: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Implementation based on PyTorch ignite package, it can support more kinds of backends.
+
+        """
+        data = data.unsqueeze(0) if ndims == 0 else data
+        # make sure the data is evenly-divisible on multi-GPUs
+        all_lens: List[int] = idist.all_gather(length)
+        max_len: int = max(all_lens)
+        if length < max_len:
+            size = [max_len - length] + list(data.shape[1:])
+            data = torch.cat([data, data.new_full(size, 0)], dim=0)
+        # all gather across all processes
+        output = idist.all_gather(data)
+        # delete the padding NaN items
+        if ndims == 0:
+            # if all the input data doesn't have batch dim, unbind to a list of 0-dim Tensors
+            return list(torch.unbind(output, dim=0))
+        return [output[i * max_len : i * max_len + l, ...] for i, l in enumerate(all_lens)]
+
+    output: List[torch.Tensor]
+    if has_ignite:
+        if idist.get_world_size() <= 1:
+            return data
+        output = _ignite_all_gather(data=data)
+    elif dist.is_available() and dist.is_initialized():
+        if dist.get_world_size() <= 1:
+            return data
+        output = _torch_all_gather(data=data)
+    else:
+        return data
+
+    return torch.cat(output, dim=0) if concat else output
+
+
+def string_list_all_gather(strings: List[str], delimiter: str = "\t") -> List[str]:
+    """
+    Utility function for distributed data parallel to all gather a list of strings.
+    Refer to the idea of ignite `all_gather(string)`:
+    https://github.com/pytorch/ignite/blob/master/ignite/distributed/utils.py#L346.
+
+    Note: If has ignite installed, will execute based on ignite distributed APIs, otherwise, if the native
+    PyTorch distributed group initialized, will execute based on native PyTorch distributed APIs.
+
+    Args:
+        strings: a list of strings to all gather.
+        delimiter: use the delimiter to join the string list to be a long string,
+            then all gather across ranks and split to a list. default to "\t".
+
+    """
+    world_size: int = 1
+    if has_ignite:
+        world_size = idist.get_world_size()
+    elif dist.is_available() and dist.is_initialized():
+        world_size = dist.get_world_size()
+
+    if world_size <= 1:
+        return strings
+
+    joined = delimiter.join(strings)
+    gathered = evenly_divisible_all_gather(torch.tensor(bytearray(joined, "utf-8"), dtype=torch.long), concat=False)
+    gathered = [bytearray(g.tolist()).decode("utf-8").split(delimiter) for g in gathered]
+
+    return [i for k in gathered for i in k]
diff --git a/monai/utils/misc.py b/monai/utils/misc.py
--- a/monai/utils/misc.py
+++ b/monai/utils/misc.py
@@ -17,11 +17,10 @@
 import warnings
 from ast import literal_eval
 from distutils.util import strtobool
-from typing import Any, Callable, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 import torch
-import torch.distributed as dist
 
 __all__ = [
     "zip_with",
@@ -42,9 +41,6 @@
     "dtype_numpy_to_torch",
     "MAX_SEED",
     "copy_to_device",
-    "get_dist_device",
-    "evenly_divisible_all_gather",
-    "string_list_all_gather",
     "ImageMetaKey",
 ]
 
@@ -356,95 +352,6 @@ def copy_to_device(
     return obj
 
 
-def get_dist_device():
-    """
-    Get the expected target device in the distributed data parallel.
-    For NCCL backend, return GPU device of current process.
-    For GLOO backend, return CPU.
-    For any other backends, return None as the default, tensor.to(None) will not change the device.
-
-    """
-    if dist.is_initialized():
-        backend = dist.get_backend()
-        if backend == "nccl" and torch.cuda.is_available():
-            return torch.device(f"cuda:{torch.cuda.current_device()}")
-        elif backend == "gloo":
-            return torch.device("cpu")
-    return None
-
-
-def evenly_divisible_all_gather(data: torch.Tensor, concat: bool = True):
-    """
-    Utility function for distributed data parallel to pad at first dim to make it evenly divisible and all_gather.
-    The input data of every rank should have the same number of dimensions, only the first dim can be different.
-
-    Args:
-        data: source tensor to pad and execute all_gather in distributed data parallel.
-        concat: whether to concat the gathered list to be a Tensor, if False, return a list
-            of Tensors, similar behavior as torch.distributed.all_gather(). default to True.
-
-    Note:
-        The input data on different ranks must have exactly same `dtype`.
-
-    """
-    if not isinstance(data, torch.Tensor):
-        raise ValueError("input data must be PyTorch Tensor.")
-
-    world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
-    if world_size <= 1:
-        return data
-
-    device = get_dist_device()
-    orig_device = data.device
-    data = data.to(device)
-    # data of all the ranks must have same number of dimensions
-    ndims = data.ndimension()
-    if ndims == 0:
-        # tensor must have batch dimension
-        data = data.unsqueeze(0)
-    # make sure the data is evenly-divisible on multi-GPUs
-    length: int = data.shape[0]
-    length_tensor = torch.as_tensor([length], device=device)
-    all_lens = [torch.zeros_like(length_tensor) for _ in range(world_size)]
-    dist.all_gather(all_lens, length_tensor)
-    all_lens_: List[int] = [int(i.item()) for i in all_lens]
-
-    max_len: int = max(all_lens_)
-    if length < max_len:
-        size = [max_len - length] + list(data.shape[1:])
-        data = torch.cat([data, data.new_full(size, 0)], dim=0)
-    # all gather across all processes
-    output = [torch.zeros_like(data) for _ in range(world_size)]
-    dist.all_gather(output, data)
-    # remove the padding items, if all the input data doesn't have batch dim, suqeeze the first dim
-    output = [(o.squeeze(0) if ndims == 0 else o[:l, ...]).to(orig_device) for o, l in zip(output, all_lens_)]
-
-    return torch.cat(output, dim=0) if concat else output
-
-
-def string_list_all_gather(strings: List[str], delimiter: str = "\t") -> List[str]:
-    """
-    Utility function for distributed data parallel to all gather a list of strings.
-    Refer to the idea of ignite `all_gather(string)`:
-    https://github.com/pytorch/ignite/blob/master/ignite/distributed/utils.py#L346.
-
-    Args:
-        strings: a list of strings to all gather.
-        delimiter: use the delimiter to join the string list to be a long string,
-            then all gather across ranks and split to a list. default to "\t".
-
-    """
-    world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
-    if world_size <= 1:
-        return strings
-
-    joined = delimiter.join(strings)
-    gathered = evenly_divisible_all_gather(torch.tensor(bytearray(joined, "utf-8"), dtype=torch.long), concat=False)
-    gathered = [bytearray(g.tolist()).decode("utf-8").split(delimiter) for g in gathered]
-
-    return [i for k in gathered for i in k]
-
-
 class ImageMetaKey:
     """
     Common key names in the meta data header of images


Test patch:
diff --git a/tests/test_handler_checkpoint_saver.py b/tests/test_handler_checkpoint_saver.py
--- a/tests/test_handler_checkpoint_saver.py
+++ b/tests/test_handler_checkpoint_saver.py
@@ -31,7 +31,7 @@
     None,
     False,
     False,
-    1.0,
+    False,
     True,
     0,
     None,
@@ -47,7 +47,7 @@
     None,
     False,
     True,
-    1,
+    False,
     False,
     0,
     None,
@@ -63,7 +63,7 @@
     None,
     False,
     True,
-    1.0,
+    False,
     True,
     2,
     2,
@@ -79,7 +79,7 @@
     None,
     False,
     False,
-    1.0,
+    False,
     False,
     10,
     2,
@@ -95,7 +95,7 @@
     None,
     False,
     False,
-    1.0,
+    False,
     True,
     0,
     None,
@@ -103,11 +103,11 @@
     True,
 ]
 
-TEST_CASE_6 = [True, "final_model.pt", False, None, 1, None, False, False, 1.0, True, 0, None, ["final_model.pt"]]
+TEST_CASE_6 = [True, "final_model.pt", False, None, 1, None, False, False, False, True, 0, None, ["final_model.pt"]]
 
-TEST_CASE_7 = [False, None, True, "val_loss", 1, "model.pt", False, False, 1.0, True, 0, None, ["model.pt"]]
+TEST_CASE_7 = [False, None, True, "val_loss", 1, "model.pt", False, False, False, True, 0, None, ["model.pt"]]
 
-TEST_CASE_8 = [False, None, True, "val_loss", 1, "model.pt", False, True, 1.0, True, 0, None, ["model.pt"]]
+TEST_CASE_8 = [False, None, True, "val_loss", 1, "model.pt", False, True, False, True, 0, None, ["model.pt"]]
 
 
 class TestHandlerCheckpointSaver(unittest.TestCase):
@@ -133,7 +133,7 @@ def test_file(
         key_metric_filename,
         key_metric_save_state,
         key_metric_greater_or_equal,
-        key_metric_score_sign,
+        key_metric_negative_sign,
         epoch_level,
         save_interval,
         n_saved,
@@ -168,7 +168,7 @@ def _train_func(engine, batch):
                 key_metric_filename,
                 key_metric_save_state,
                 key_metric_greater_or_equal,
-                key_metric_score_sign,
+                key_metric_negative_sign,
                 epoch_level,
                 save_interval,
                 n_saved,
@@ -218,7 +218,7 @@ def _train_func(engine, batch):
                 key_metric_name="val_loss",
                 key_metric_n_saved=2,
                 key_metric_save_state=True,
-                key_metric_score_sign=-1,
+                key_metric_negative_sign=True,
             ).attach(engine)
             engine.run(range(3), max_epochs=3)
 
@@ -228,7 +228,7 @@ def _train_func(engine, batch):
                 save_key_metric=True,
                 key_metric_name="val_loss",
                 key_metric_n_saved=2,
-                key_metric_score_sign=-1,
+                key_metric_negative_sign=True,
             )
             engine = Engine(_train_func)
             CheckpointLoader(os.path.join(tempdir, "net_key_metric=-6.pt"), {"checkpointer": saver}).attach(engine)