facebook/StyleNeRF

1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved2 3import logging4import os5import pickle6import random7import socket8import struct9import subprocess10import warnings11import tempfile12import uuid13 14 15from datetime import date16from pathlib import Path17from collections import OrderedDict18from typing import Any, Dict, Mapping19 20import torch21import torch.distributed as dist22 23 24logger = logging.getLogger(__name__)25 26 27def is_master(args):28    return args.distributed_rank == 029 30 31def init_distributed_mode(rank, args):32    if "WORLD_SIZE" in os.environ:33        args.world_size = int(os.environ["WORLD_SIZE"])34    35    if args.launcher == 'spawn':  # single node with multiprocessing.spawn36        args.world_size = args.num_gpus37        args.rank = rank38        args.gpu = rank39    40    elif 'RANK' in os.environ:41        args.rank = int(os.environ["RANK"])42        args.gpu = int(os.environ['LOCAL_RANK'])43    44    elif 'SLURM_PROCID' in os.environ:45        args.rank = int(os.environ['SLURM_PROCID'])46        args.gpu = args.rank % torch.cuda.device_count()47    48    if args.world_size == 1:49        return50 51    if 'MASTER_ADDR' in os.environ:52        args.dist_url = 'tcp://{}:{}'.format(os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])53 54    print(f'gpu={args.gpu}, rank={args.rank}, world_size={args.world_size}')55    args.distributed = True56    torch.cuda.set_device(args.gpu)57    args.dist_backend = 'nccl'58    print('| distributed init (rank {}): {}'.format(args.rank, args.dist_url), flush=True)59    60    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,61                                         world_size=args.world_size, rank=args.rank)62    torch.distributed.barrier()63 64 65def gather_list_and_concat(tensor):66    gather_t = [torch.ones_like(tensor) for _ in range(dist.get_world_size())]67    dist.all_gather(gather_t, tensor)68    return torch.cat(gather_t)69 70 71def get_rank():72    return dist.get_rank()73 74 75def get_world_size():76    return dist.get_world_size()77 78 79def get_default_group():80    return dist.group.WORLD81 82 83def all_gather_list(data, group=None, max_size=16384):84    """Gathers arbitrary data from all nodes into a list.85 86    Similar to :func:`~torch.distributed.all_gather` but for arbitrary Python87    data. Note that *data* must be picklable.88 89    Args:90        data (Any): data from the local worker to be gathered on other workers91        group (optional): group of the collective92        max_size (int, optional): maximum size of the data to be gathered93            across workers94    """95    rank = get_rank()96    world_size = get_world_size()97 98    buffer_size = max_size * world_size99    if not hasattr(all_gather_list, '_buffer') or \100            all_gather_list._buffer.numel() < buffer_size:101        all_gather_list._buffer = torch.cuda.ByteTensor(buffer_size)102        all_gather_list._cpu_buffer = torch.ByteTensor(max_size).pin_memory()103    buffer = all_gather_list._buffer104    buffer.zero_()105    cpu_buffer = all_gather_list._cpu_buffer106 107    data = data.cpu()108    enc = pickle.dumps(data)109    enc_size = len(enc)110    header_size = 4  # size of header that contains the length of the encoded data111    size = header_size + enc_size112    if size > max_size:113        raise ValueError('encoded data size ({}) exceeds max_size ({})'.format(size, max_size))114 115    header = struct.pack(">I", enc_size)116    cpu_buffer[:size] = torch.ByteTensor(list(header + enc))117    start = rank * max_size118    buffer[start:start + size].copy_(cpu_buffer[:size])119 120    all_reduce(buffer, group=group)121 122    buffer = buffer.cpu()123    try:124        result = []125        for i in range(world_size):126            out_buffer = buffer[i * max_size:(i + 1) * max_size]127            enc_size, = struct.unpack(">I", bytes(out_buffer[:header_size].tolist()))128            if enc_size > 0:129                result.append(pickle.loads(bytes(out_buffer[header_size:header_size + enc_size].tolist())))130        return result131    except pickle.UnpicklingError:132        raise Exception(133            'Unable to unpickle data from other workers. all_gather_list requires all '134            'workers to enter the function together, so this error usually indicates '135            'that the workers have fallen out of sync somehow. Workers can fall out of '136            'sync if one of them runs out of memory, or if there are other conditions '137            'in your training script that can cause one worker to finish an epoch '138            'while other workers are still iterating over their portions of the data. '139            'Try rerunning with --ddp-backend=no_c10d and see if that helps.'140        )141 142 143def all_reduce_dict(144    data: Mapping[str, Any],145    device,146    group=None,147) -> Dict[str, Any]:148    """149    AllReduce a dictionary of values across workers. We separately150    reduce items that are already on the device and items on CPU for151    better performance.152 153    Args:154        data (Mapping[str, Any]): dictionary of data to all-reduce, but155            cannot be a nested dictionary156        device (torch.device): device for the reduction157        group (optional): group of the collective158    """159    data_keys = list(data.keys())160 161    # We want to separately reduce items that are already on the162    # device and items on CPU for performance reasons.163    cpu_data = OrderedDict()164    device_data = OrderedDict()165    for k in data_keys:166        t = data[k]167        if not torch.is_tensor(t):168            cpu_data[k] = torch.tensor(t, dtype=torch.double)169        elif t.device.type != device.type:170            cpu_data[k] = t.to(dtype=torch.double)171        else:172            device_data[k] = t.to(dtype=torch.double)173 174    def _all_reduce_dict(data: OrderedDict):175        if len(data) == 0:176            return data177        buf = torch.stack(list(data.values())).to(device=device)178        all_reduce(buf, group=group)179        return {k: buf[i] for i, k in enumerate(data)}180 181    cpu_data = _all_reduce_dict(cpu_data)182    device_data = _all_reduce_dict(device_data)183 184    def get_from_stack(key):185        if key in cpu_data:186            return cpu_data[key]187        elif key in device_data:188            return device_data[key]189        raise KeyError190 191    return OrderedDict([(key, get_from_stack(key)) for key in data_keys])192 193 194def get_shared_folder() -> Path:195    user = os.getenv("USER")196    if Path("/checkpoint/").is_dir():197        p = Path(f"/checkpoint/{user}/experiments")198        p.mkdir(exist_ok=True)199        return p200    else:201        p = Path(f"/tmp/experiments")202        p.mkdir(exist_ok=True)203        return p204 205 206def get_init_file():207    # Init file must not exist, but it's parent dir must exist.208    os.makedirs(str(get_shared_folder()), exist_ok=True)209    init_file = Path(str(get_shared_folder()) + f"/{uuid.uuid4().hex}_init")210    if init_file.exists():211        os.remove(str(init_file))212    return init_file213 214