pyanomaly.datatools.dataclass.sampler package

Submodules

pyanomaly.datatools.dataclass.sampler.common module

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

pyanomaly.datatools.dataclass.sampler.common.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

pyanomaly.datatools.dataclass.sampler.common.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.

pyanomaly.datatools.dataclass.sampler.common.get_local_rank() → int
Returns:

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

pyanomaly.datatools.dataclass.sampler.common.get_local_size() → int
Returns:

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

pyanomaly.datatools.dataclass.sampler.common.get_rank() → int
pyanomaly.datatools.dataclass.sampler.common.get_world_size() → int
pyanomaly.datatools.dataclass.sampler.common.is_main_process() → bool
pyanomaly.datatools.dataclass.sampler.common.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.

pyanomaly.datatools.dataclass.sampler.common.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.

pyanomaly.datatools.dataclass.sampler.common.synchronize()

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

pyanomaly.datatools.dataclass.sampler.dist_inf_sampler module

class pyanomaly.datatools.dataclass.sampler.dist_inf_sampler.DistTrainSampler(size, shuffle=True, seed=None, start=0)

Bases: torch.utils.data.sampler.Sampler

refer https://github.com/facebookresearch/detectron2/blob/master/detectron2/data/samplers/distributed_sampler.py

In training, we only care about the "infinite stream" of training data. So this sampler produces an infinite stream of indices and all workers cooperate to correctly shuffle the indices and sample different indices. The samplers in each worker effectively produces indices[worker_id::num_workers] where indices is an infinite stream of indices consisting of shuffle(range(size)) + shuffle(range(size)) + ... (if shuffle is True) or range(size) + range(size) + ... (if shuffle is False)

class pyanomaly.datatools.dataclass.sampler.dist_inf_sampler.InferenceSampler(size, start=0)

Bases: torch.utils.data.sampler.Sampler

refer https://github.com/facebookresearch/detectron2/blob/master/detectron2/data/samplers/distributed_sampler.py Produce indices for inference. Inference needs to run on the __exact__ set of samples, therefore when the total number of samples is not divisible by the number of workers, this sampler produces different number of samples on different workers.

pyanomaly.datatools.dataclass.sampler.inf_sampler module

class pyanomaly.datatools.dataclass.sampler.inf_sampler.TrainSampler(size, seed, start=0, shuffle=True)

Bases: torch.utils.data.sampler.Sampler

Generate the infinite stream, not support dist However, the inference not use this!!

Module contents

@author: Yuhao Cheng @contact: yuhao.cheng[at]outlook.com