sczhou/CodeFormer

1# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E5012import functools3import os4import subprocess5import torch6import torch.distributed as dist7import torch.multiprocessing as mp8 9 10def init_dist(launcher, backend='nccl', **kwargs):11    if mp.get_start_method(allow_none=True) is None:12        mp.set_start_method('spawn')13    if launcher == 'pytorch':14        _init_dist_pytorch(backend, **kwargs)15    elif launcher == 'slurm':16        _init_dist_slurm(backend, **kwargs)17    else:18        raise ValueError(f'Invalid launcher type: {launcher}')19 20 21def _init_dist_pytorch(backend, **kwargs):22    rank = int(os.environ['RANK'])23    num_gpus = torch.cuda.device_count()24    torch.cuda.set_device(rank % num_gpus)25    dist.init_process_group(backend=backend, **kwargs)26 27 28def _init_dist_slurm(backend, port=None):29    """Initialize slurm distributed training environment.30 31    If argument ``port`` is not specified, then the master port will be system32    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system33    environment variable, then a default port ``29500`` will be used.34 35    Args:36        backend (str): Backend of torch.distributed.37        port (int, optional): Master port. Defaults to None.38    """39    proc_id = int(os.environ['SLURM_PROCID'])40    ntasks = int(os.environ['SLURM_NTASKS'])41    node_list = os.environ['SLURM_NODELIST']42    num_gpus = torch.cuda.device_count()43    torch.cuda.set_device(proc_id % num_gpus)44    addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')45    # specify master port46    if port is not None:47        os.environ['MASTER_PORT'] = str(port)48    elif 'MASTER_PORT' in os.environ:49        pass  # use MASTER_PORT in the environment variable50    else:51        # 29500 is torch.distributed default port52        os.environ['MASTER_PORT'] = '29500'53    os.environ['MASTER_ADDR'] = addr54    os.environ['WORLD_SIZE'] = str(ntasks)55    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)56    os.environ['RANK'] = str(proc_id)57    dist.init_process_group(backend=backend)58 59 60def get_dist_info():61    if dist.is_available():62        initialized = dist.is_initialized()63    else:64        initialized = False65    if initialized:66        rank = dist.get_rank()67        world_size = dist.get_world_size()68    else:69        rank = 070        world_size = 171    return rank, world_size72 73 74def master_only(func):75 76    @functools.wraps(func)77    def wrapper(*args, **kwargs):78        rank, _ = get_dist_info()79        if rank == 0:80            return func(*args, **kwargs)81 82    return wrapper83