mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
# Copyright 2022 The Plenoptix Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""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()
|