mirror of https://github.com/open-mmlab/mmcv.git
79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
# Copyright (c) Open-MMLab. All rights reserved.
|
|
import functools
|
|
import os
|
|
import subprocess
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
|
|
from mmcv.utils import TORCH_VERSION
|
|
|
|
|
|
def init_dist(launcher, backend='nccl', **kwargs):
|
|
if mp.get_start_method(allow_none=True) is None:
|
|
mp.set_start_method('spawn')
|
|
if launcher == 'pytorch':
|
|
_init_dist_pytorch(backend, **kwargs)
|
|
elif launcher == 'mpi':
|
|
_init_dist_mpi(backend, **kwargs)
|
|
elif launcher == 'slurm':
|
|
_init_dist_slurm(backend, **kwargs)
|
|
else:
|
|
raise ValueError(f'Invalid launcher type: {launcher}')
|
|
|
|
|
|
def _init_dist_pytorch(backend, **kwargs):
|
|
# TODO: use local_rank instead of rank % num_gpus
|
|
rank = int(os.environ['RANK'])
|
|
num_gpus = torch.cuda.device_count()
|
|
torch.cuda.set_device(rank % num_gpus)
|
|
dist.init_process_group(backend=backend, **kwargs)
|
|
|
|
|
|
def _init_dist_mpi(backend, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
def _init_dist_slurm(backend, port=29500):
|
|
proc_id = int(os.environ['SLURM_PROCID'])
|
|
ntasks = int(os.environ['SLURM_NTASKS'])
|
|
node_list = os.environ['SLURM_NODELIST']
|
|
num_gpus = torch.cuda.device_count()
|
|
torch.cuda.set_device(proc_id % num_gpus)
|
|
addr = subprocess.getoutput(
|
|
f'scontrol show hostname {node_list} | head -n1')
|
|
os.environ['MASTER_PORT'] = str(port)
|
|
os.environ['MASTER_ADDR'] = addr
|
|
os.environ['WORLD_SIZE'] = str(ntasks)
|
|
os.environ['RANK'] = str(proc_id)
|
|
dist.init_process_group(backend=backend)
|
|
|
|
|
|
def get_dist_info():
|
|
if TORCH_VERSION < '1.0':
|
|
initialized = dist._initialized
|
|
else:
|
|
if dist.is_available():
|
|
initialized = dist.is_initialized()
|
|
else:
|
|
initialized = False
|
|
if initialized:
|
|
rank = dist.get_rank()
|
|
world_size = dist.get_world_size()
|
|
else:
|
|
rank = 0
|
|
world_size = 1
|
|
return rank, world_size
|
|
|
|
|
|
def master_only(func):
|
|
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
rank, _ = get_dist_info()
|
|
if rank == 0:
|
|
return func(*args, **kwargs)
|
|
|
|
return wrapper
|