mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
"""functionality to handle multiprocessing syncing and communicating"""
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
_LOCAL_PROCESS_GROUP = None
|
|
|
|
|
|
def get_world_size() -> int:
|
|
"""Get total number of available gpus"""
|
|
if not dist.is_available():
|
|
return 1
|
|
if not dist.is_initialized():
|
|
return 1
|
|
return dist.get_world_size()
|
|
|
|
|
|
def get_rank() -> int:
|
|
"""Get global rank of current thread"""
|
|
if not dist.is_available():
|
|
return 0
|
|
if not dist.is_initialized():
|
|
return 0
|
|
return dist.get_rank()
|
|
|
|
|
|
def get_local_rank() -> int:
|
|
"""The rank of the current process within the local (per-machine) process group."""
|
|
if not dist.is_available():
|
|
return 0
|
|
if not dist.is_initialized():
|
|
return 0
|
|
assert (
|
|
_LOCAL_PROCESS_GROUP is not None
|
|
), "Local process group is not created! Please use launch() to spawn processes!"
|
|
return dist.get_rank(group=_LOCAL_PROCESS_GROUP)
|
|
|
|
|
|
def get_local_size() -> int:
|
|
"""
|
|
The size of the per-machine process group,
|
|
i.e. the number of processes per machine.
|
|
"""
|
|
if not dist.is_available():
|
|
return 1
|
|
if not dist.is_initialized():
|
|
return 1
|
|
return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)
|
|
|
|
|
|
def is_main_process() -> bool:
|
|
"""check to see if you are currently on the main process"""
|
|
return get_rank() == 0
|
|
|
|
|
|
def synchronize(world_size):
|
|
"""
|
|
Helper function to synchronize (barrier) among all processes when
|
|
using distributed training
|
|
"""
|
|
if not dist.is_available():
|
|
return
|
|
if not dist.is_initialized():
|
|
return
|
|
world_size = dist.get_world_size()
|
|
if world_size == 1:
|
|
return
|
|
if dist.get_backend() == dist.Backend.NCCL:
|
|
# This argument is needed to avoid warnings.
|
|
# It's valid only for NCCL backend.
|
|
dist.barrier(device_ids=[torch.cuda.current_device()])
|
|
else:
|
|
dist.barrier()
|