fastreid.utils.comm Namespace Reference

Functions
int	get_world_size ()
int	get_rank ()
int	get_local_rank ()
int	get_local_size ()
bool	is_main_process ()
	synchronize ()
	_get_global_gloo_group ()
	_serialize_to_tensor (data, group)
	_pad_to_largest_tensor (tensor, group)
	all_gather (data, group=None)
	gather (data, dst=0, group=None)
	shared_random_seed ()
	reduce_dict (input_dict, average=True)

Variables
	_LOCAL_PROCESS_GROUP = None

Detailed Description

This file contains primitives for multi-gpu communication.
This is useful when doing distributed training.

Function Documentation

_get_global_gloo_group()

fastreid.utils.comm._get_global_gloo_group ( )

protected

Return a process group based on gloo backend, containing all the ranks
The result is cached.

Definition at line 82 of file comm.py.

def _get_global_gloo_group():
    """
    Return a process group based on gloo backend, containing all the ranks
    The result is cached.
    """
    if dist.get_backend() == "nccl":
        return dist.new_group(backend="gloo")
    else:
        return dist.group.WORLD
 
 

_pad_to_largest_tensor()

fastreid.utils.comm._pad_to_largest_tensor ( tensor, group )

protected

Returns:
    list[int]: size of the tensor, on each rank
    Tensor: padded tensor that has the max size

Definition at line 111 of file comm.py.

def _pad_to_largest_tensor(tensor, group):
    """
    Returns:
        list[int]: size of the tensor, on each rank
        Tensor: padded tensor that has the max size
    """
    world_size = dist.get_world_size(group=group)
    assert (
            world_size >= 1
    ), "comm.gather/all_gather must be called from ranks within the given group!"
    local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device)
    size_list = [
        torch.zeros([1], dtype=torch.int64, device=tensor.device) for _ in range(world_size)
    ]
    dist.all_gather(size_list, local_size, group=group)
    size_list = [int(size.item()) for size in size_list]
 
    max_size = max(size_list)
 
    # we pad the tensor because torch all_gather does not support
    # gathering tensors of different shapes
    if local_size != max_size:
        padding = torch.zeros((max_size - local_size,), dtype=torch.uint8, device=tensor.device)
        tensor = torch.cat((tensor, padding), dim=0)
    return size_list, tensor
 
 

_serialize_to_tensor()

fastreid.utils.comm._serialize_to_tensor ( data, group )

protected

Definition at line 93 of file comm.py.

def _serialize_to_tensor(data, group):
    backend = dist.get_backend(group)
    assert backend in ["gloo", "nccl"]
    device = torch.device("cpu" if backend == "gloo" else "cuda")
 
    buffer = pickle.dumps(data)
    if len(buffer) > 1024 ** 3:
        logger = logging.getLogger(__name__)
        logger.warning(
            "Rank {} trying to all-gather {:.2f} GB of data on device {}".format(
                get_rank(), len(buffer) / (1024 ** 3), device
            )
        )
    storage = torch.ByteStorage.from_buffer(buffer)
    tensor = torch.ByteTensor(storage).to(device=device)
    return tensor
 
 

all_gather()

fastreid.utils.comm.all_gather	(		data,
			group = None )

Run all_gather on arbitrary picklable data (not necessarily tensors).
Args:
    data: any picklable object
    group: a torch process group. By default, will use a group which
        contains all ranks on gloo backend.
Returns:
    list[data]: list of data gathered from each rank

Definition at line 138 of file comm.py.

def all_gather(data, group=None):
    """
    Run all_gather on arbitrary picklable data (not necessarily tensors).
    Args:
        data: any picklable object
        group: a torch process group. By default, will use a group which
            contains all ranks on gloo backend.
    Returns:
        list[data]: list of data gathered from each rank
    """
    if get_world_size() == 1:
        return [data]
    if group is None:
        group = _get_global_gloo_group()
    if dist.get_world_size(group) == 1:
        return [data]
 
    tensor = _serialize_to_tensor(data, group)
 
    size_list, tensor = _pad_to_largest_tensor(tensor, group)
    max_size = max(size_list)
 
    # receiving Tensor from all ranks
    tensor_list = [
        torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
    ]
    dist.all_gather(tensor_list, tensor, group=group)
 
    data_list = []
    for size, tensor in zip(size_list, tensor_list):
        buffer = tensor.cpu().numpy().tobytes()[:size]
        data_list.append(pickle.loads(buffer))
 
    return data_list
 
 

gather()

fastreid.utils.comm.gather	(		data,
			dst = 0,
			group = None )

Run gather on arbitrary picklable data (not necessarily tensors).
Args:
    data: any picklable object
    dst (int): destination rank
    group: a torch process group. By default, will use a group which
        contains all ranks on gloo backend.
Returns:
    list[data]: on dst, a list of data gathered from each rank. Otherwise,
        an empty list.

Definition at line 174 of file comm.py.

def gather(data, dst=0, group=None):
    """
    Run gather on arbitrary picklable data (not necessarily tensors).
    Args:
        data: any picklable object
        dst (int): destination rank
        group: a torch process group. By default, will use a group which
            contains all ranks on gloo backend.
    Returns:
        list[data]: on dst, a list of data gathered from each rank. Otherwise,
            an empty list.
    """
    if get_world_size() == 1:
        return [data]
    if group is None:
        group = _get_global_gloo_group()
    if dist.get_world_size(group=group) == 1:
        return [data]
    rank = dist.get_rank(group=group)
 
    tensor = _serialize_to_tensor(data, group)
    size_list, tensor = _pad_to_largest_tensor(tensor, group)
 
    # receiving Tensor from all ranks
    if rank == dst:
        max_size = max(size_list)
        tensor_list = [
            torch.empty((max_size,), dtype=torch.uint8, device=tensor.device) for _ in size_list
        ]
        dist.gather(tensor, tensor_list, dst=dst, group=group)
 
        data_list = []
        for size, tensor in zip(size_list, tensor_list):
            buffer = tensor.cpu().numpy().tobytes()[:size]
            data_list.append(pickle.loads(buffer))
        return data_list
    else:
        dist.gather(tensor, [], dst=dst, group=group)
        return []
 
 

get_local_rank()

int fastreid.utils.comm.get_local_rank

(

)

Returns:
    The rank of the current process within the local (per-machine) process group.

Definition at line 36 of file comm.py.

def get_local_rank() -> int:
    """
    Returns:
        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
    return dist.get_rank(group=_LOCAL_PROCESS_GROUP)
 
 

get_local_size()

int fastreid.utils.comm.get_local_size

(

)

Returns:
    The size of the per-machine process group,
    i.e. the number of processes per machine.

Definition at line 49 of file comm.py.

def get_local_size() -> int:
    """
    Returns:
        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)
 
 

get_rank()

int fastreid.utils.comm.get_rank

(

)

Definition at line 28 of file comm.py.

def get_rank() -> int:
    if not dist.is_available():
        return 0
    if not dist.is_initialized():
        return 0
    return dist.get_rank()
 
 

get_world_size()

int fastreid.utils.comm.get_world_size

(

)

Definition at line 20 of file comm.py.

def get_world_size() -> int:
    if not dist.is_available():
        return 1
    if not dist.is_initialized():
        return 1
    return dist.get_world_size()
 
 

is_main_process()

bool fastreid.utils.comm.is_main_process

(

)

Definition at line 62 of file comm.py.

def is_main_process() -> bool:
    return get_rank() == 0
 
 

reduce_dict()

fastreid.utils.comm.reduce_dict	(		input_dict,
			average = True )

Reduce the values in the dictionary from all processes so that process with rank
0 has the reduced results.
Args:
    input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
    average (bool): whether to do average or sum
Returns:
    a dict with the same keys as input_dict, after reduction.

Definition at line 228 of file comm.py.

def reduce_dict(input_dict, average=True):
    """
    Reduce the values in the dictionary from all processes so that process with rank
    0 has the reduced results.
    Args:
        input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
        average (bool): whether to do average or sum
    Returns:
        a dict with the same keys as input_dict, after reduction.
    """
    world_size = get_world_size()
    if world_size < 2:
        return input_dict
    with torch.no_grad():
        names = []
        values = []
        # sort the keys so that they are consistent across processes
        for k in sorted(input_dict.keys()):
            names.append(k)
            values.append(input_dict[k])
        values = torch.stack(values, dim=0)
        dist.reduce(values, dst=0)
        if dist.get_rank() == 0 and average:
            # only main process gets accumulated, so only divide by
            # world_size in this case
            values /= world_size
        reduced_dict = {k: v for k, v in zip(names, values)}
    return reduced_dict

shared_random_seed()

fastreid.utils.comm.shared_random_seed

(

)

Returns:
    int: a random number that is the same across all workers.
        If workers need a shared RNG, they can use this shared seed to
        create one.
All workers must call this function, otherwise it will deadlock.

Definition at line 215 of file comm.py.

def shared_random_seed():
    """
    Returns:
        int: a random number that is the same across all workers.
            If workers need a shared RNG, they can use this shared seed to
            create one.
    All workers must call this function, otherwise it will deadlock.
    """
    ints = np.random.randint(2 ** 31)
    all_ints = all_gather(ints)
    return all_ints[0]
 
 

synchronize()

fastreid.utils.comm.synchronize

(

)

Helper function to synchronize (barrier) among all processes when
using distributed training

Definition at line 66 of file comm.py.

def synchronize():
    """
    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
    dist.barrier()
 
 
@functools.lru_cache()

Variable Documentation

_LOCAL_PROCESS_GROUP

fastreid.utils.comm._LOCAL_PROCESS_GROUP = None

protected

Definition at line 13 of file comm.py.