feat(mlu): Support PyTorch backend on MLU. (#1515)

* feat(mlu): Support PyTorch backend on MLU.

* fix redundant device variable.

* Update mmseg/apis/train.py

Co-authored-by: Miao Zheng <76149310+MeowZheng@users.noreply.github.com>

* Update comments.

* Update mmseg/apis/train.py

* Update is_mlu_available flag.

* align util_distribution.py to mmdet.

* align util_distribution.py to mmdet.

* add build_dp, build_ddp testcase.

* Update mmseg/utils/util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update mmseg/utils/util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update mmseg/utils/util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update tests/test_utils/test_util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update tests/test_utils/test_util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update tests/test_utils/test_util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* Update tests/test_utils/test_util_distribution.py

Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

* add mmcv version check for mlu device.

Co-authored-by: Miao Zheng <76149310+MeowZheng@users.noreply.github.com>
Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>

mmseg/apis/train.py

@@ -1,20 +1,21 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import os
 import random
 import warnings
 
 import mmcv
 import numpy as np
 import torch
 import torch.distributed as dist
-from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
 from mmcv.runner import (HOOKS, DistSamplerSeedHook, EpochBasedRunner,
  build_runner, get_dist_info)
 from mmcv.utils import build_from_cfg
 
 from mmseg import digit_version
 from mmseg.core import DistEvalHook, EvalHook, build_optimizer
 from mmseg.datasets import build_dataloader, build_dataset
-from mmseg.utils import find_latest_checkpoint, get_root_logger
+from mmseg.utils import (build_ddp, build_dp, find_latest_checkpoint,
+ get_root_logger)
 
 
 def init_random_seed(seed=None, device='cuda'):
@@ -99,21 +100,23 @@ def train_segmentor(model,
  train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})}
  data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]
 
- # put model on gpus
+ # put model on devices
  if distributed:
  find_unused_parameters = cfg.get('find_unused_parameters', False)
  # Sets the `find_unused_parameters` parameter in
- # torch.nn.parallel.DistributedDataParallel
- model = MMDistributedDataParallel(
- model.cuda(),
- device_ids=[torch.cuda.current_device()],
+ # DDP wrapper
+ model = build_ddp(
+ model,
+ cfg.device,
+ device_ids=[int(os.environ['LOCAL_RANK'])],
  broadcast_buffers=False,
  find_unused_parameters=find_unused_parameters)
  else:
  if not torch.cuda.is_available():
  assert digit_version(mmcv.__version__) >= digit_version('1.4.4'), \
  'Please use MMCV >= 1.4.4 for CPU training!'
- model = MMDataParallel(model, device_ids=cfg.gpu_ids)
+ model = build_dp(model, cfg.device, device_ids=cfg.gpu_ids)
+
  # build runner
  optimizer = build_optimizer(model, cfg.optimizer)
 

mmseg/datasets/samplers/distributed_sampler.py

@@ -7,6 +7,7 @@
 from torch.utils.data import DistributedSampler as _DistributedSampler
 
 from mmseg.core.utils import sync_random_seed
+from mmseg.utils import get_device
 
 
 class DistributedSampler(_DistributedSampler):
@@ -41,7 +42,8 @@ def __init__(self,
  # in the same order based on the same seed. Then different ranks
  # could use different indices to select non-overlapped data from the
  # same data list.
- self.seed = sync_random_seed(seed)
+ device = get_device()
+ self.seed = sync_random_seed(seed, device)
 
  def __iter__(self) -> Iterator:
  """

mmseg/utils/__init__.py

@@ -3,8 +3,9 @@
 from .logger import get_root_logger
 from .misc import find_latest_checkpoint
 from .set_env import setup_multi_processes
+from .util_distribution import build_ddp, build_dp, get_device
 
 __all__ = [
  'get_root_logger', 'collect_env', 'find_latest_checkpoint',
- 'setup_multi_processes'
+ 'setup_multi_processes', 'build_ddp', 'build_dp', 'get_device'
 ]

mmseg/utils/util_distribution.py

@@ -0,0 +1,81 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import mmcv
+import torch
+from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
+
+from mmseg import digit_version
+
+dp_factory = {'cuda': MMDataParallel, 'cpu': MMDataParallel}
+
+ddp_factory = {'cuda': MMDistributedDataParallel}
+
+
+def build_dp(model, device='cuda', dim=0, *args, **kwargs):
+ """build DataParallel module by device type.
+
+ if device is cuda, return a MMDataParallel module; if device is mlu,
+ return a MLUDataParallel module.
+
+ Args:
+ model (:class:`nn.Module`): module to be parallelized.
+ device (str): device type, cuda, cpu or mlu. Defaults to cuda.
+ dim (int): Dimension used to scatter the data. Defaults to 0.
+
+ Returns:
+ :class:`nn.Module`: parallelized module.
+ """
+ if device == 'cuda':
+ model = model.cuda()
+ elif device == 'mlu':
+ assert digit_version(mmcv.__version__) >= digit_version('1.5.0'), \
+ 'Please use MMCV >= 1.5.0 for MLU training!'
+ from mmcv.device.mlu import MLUDataParallel
+ dp_factory['mlu'] = MLUDataParallel
+ model = model.mlu()
+
+ return dp_factory[device](model, dim=dim, *args, **kwargs)
+
+
+def build_ddp(model, device='cuda', *args, **kwargs):
+ """Build DistributedDataParallel module by device type.
+
+ If device is cuda, return a MMDistributedDataParallel module;
+ if device is mlu, return a MLUDistributedDataParallel module.
+
+ Args:
+ model (:class:`nn.Module`): module to be parallelized.
+ device (str): device type, mlu or cuda.
+
+ Returns:
+ :class:`nn.Module`: parallelized module.
+
+ References:
+ .. [1] https://pytorch.org/docs/stable/generated/torch.nn.parallel.
+ DistributedDataParallel.html
+ """
+ assert device in ['cuda', 'mlu'], 'Only available for cuda or mlu devices.'
+ if device == 'cuda':
+ model = model.cuda()
+ elif device == 'mlu':
+ assert digit_version(mmcv.__version__) >= digit_version('1.5.0'), \
+ 'Please use MMCV >= 1.5.0 for MLU training!'
+ from mmcv.device.mlu import MLUDistributedDataParallel
+ ddp_factory['mlu'] = MLUDistributedDataParallel
+ model = model.mlu()
+
+ return ddp_factory[device](model, *args, **kwargs)
+
+
+def is_mlu_available():
+ """Returns a bool indicating if MLU is currently available."""
+ return hasattr(torch, 'is_mlu_available') and torch.is_mlu_available()
+
+
+def get_device():
+ """Returns an available device, cpu, cuda or mlu."""
+ is_device_available = {
+ 'cuda': torch.cuda.is_available(),
+ 'mlu': is_mlu_available()
+ }
+ device_list = [k for k, v in is_device_available.items() if v]
+ return device_list[0] if len(device_list) == 1 else 'cpu'

tests/test_utils/test_util_distribution.py

@@ -0,0 +1,68 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from unittest.mock import MagicMock, patch
+
+import mmcv
+import torch
+import torch.nn as nn
+from mmcv.parallel import (MMDataParallel, MMDistributedDataParallel,
+ is_module_wrapper)
+
+from mmseg import digit_version
+from mmseg.utils import build_ddp, build_dp
+
+
+def mock(*args, **kwargs):
+ pass
+
+
+class Model(nn.Module):
+
+ def __init__(self):
+ super().__init__()
+ self.conv = nn.Conv2d(2, 2, 1)
+
+ def forward(self, x):
+ return self.conv(x)
+
+
+@patch('torch.distributed._broadcast_coalesced', mock)
+@patch('torch.distributed.broadcast', mock)
+@patch('torch.nn.parallel.DistributedDataParallel._ddp_init_helper', mock)
+def test_build_dp():
+ model = Model()
+ assert not is_module_wrapper(model)
+
+ mmdp = build_dp(model, 'cpu')
+ assert isinstance(mmdp, MMDataParallel)
+
+ if torch.cuda.is_available():
+ mmdp = build_dp(model, 'cuda')
+ assert isinstance(mmdp, MMDataParallel)
+
+ if digit_version(mmcv.__version__) >= digit_version('1.5.0'):
+ from mmcv.device.mlu import MLUDataParallel
+ from mmcv.utils import IS_MLU_AVAILABLE
+ if IS_MLU_AVAILABLE:
+ mludp = build_dp(model, 'mlu')
+ assert isinstance(mludp, MLUDataParallel)
+
+
+@patch('torch.distributed._broadcast_coalesced', mock)
+@patch('torch.distributed.broadcast', mock)
+@patch('torch.nn.parallel.DistributedDataParallel._ddp_init_helper', mock)
+def test_build_ddp():
+ model = Model()
+ assert not is_module_wrapper(model)
+
+ if torch.cuda.is_available():
+ mmddp = build_ddp(
+ model, 'cuda', device_id=[0], process_group=MagicMock())
+ assert isinstance(mmddp, MMDistributedDataParallel)
+
+ if digit_version(mmcv.__version__) >= digit_version('1.5.0'):
+ from mmcv.device.mlu import MLUDistributedDataParallel
+ from mmcv.utils import IS_MLU_AVAILABLE
+ if IS_MLU_AVAILABLE:
+ mluddp = build_ddp(
+ model, 'mlu', device_ids=[0], process_group=MagicMock())
+ assert isinstance(mluddp, MLUDistributedDataParallel)

tools/test.py

@@ -9,7 +9,6 @@
 import mmcv
 import torch
 from mmcv.cnn.utils import revert_sync_batchnorm
-from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
 from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
  wrap_fp16_model)
 from mmcv.utils import DictAction
@@ -18,7 +17,7 @@
 from mmseg.apis import multi_gpu_test, single_gpu_test
 from mmseg.datasets import build_dataloader, build_dataset
 from mmseg.models import build_segmentor
-from mmseg.utils import setup_multi_processes
+from mmseg.utils import build_ddp, build_dp, get_device, setup_multi_processes
 
 
 def parse_args():
@@ -260,6 +259,7 @@ def main():
  else:
  tmpdir = None
 
+ cfg.device = get_device()
  if not distributed:
  warnings.warn(
  'SyncBN is only supported with DDP. To be compatible with DP, '
@@ -269,7 +269,7 @@ def main():
  assert digit_version(mmcv.__version__) >= digit_version('1.4.4'), \
  'Please use MMCV >= 1.4.4 for CPU training!'
  model = revert_sync_batchnorm(model)
- model = MMDataParallel(model, device_ids=cfg.gpu_ids)
+ model = build_dp(model, cfg.device, device_ids=cfg.gpu_ids)
  results = single_gpu_test(
  model,
  data_loader,
@@ -281,9 +281,10 @@ def main():
  format_only=args.format_only or eval_on_format_results,
  format_args=eval_kwargs)
  else:
- model = MMDistributedDataParallel(
- model.cuda(),
- device_ids=[torch.cuda.current_device()],
+ model = build_ddp(
+ model,
+ cfg.device,
+ device_ids=[int(os.environ['LOCAL_RANK'])],
  broadcast_buffers=False)
  results = multi_gpu_test(
  model,

tools/train.py

@@ -17,7 +17,8 @@
 from mmseg.apis import init_random_seed, set_random_seed, train_segmentor
 from mmseg.datasets import build_dataset
 from mmseg.models import build_segmentor
-from mmseg.utils import collect_env, get_root_logger, setup_multi_processes
+from mmseg.utils import (collect_env, get_device, get_root_logger,
+ setup_multi_processes)
 
 
 def parse_args():
@@ -184,7 +185,8 @@ def main():
  logger.info(f'Config:\n{cfg.pretty_text}')
 
  # set random seeds
- seed = init_random_seed(args.seed)
+ cfg.device = get_device()
+ seed = init_random_seed(args.seed, device=cfg.device)
  seed = seed + dist.get_rank() if args.diff_seed else seed
  logger.info(f'Set random seed to {seed}, '
  f'deterministic: {args.deterministic}')