API Reference

torchfoo

seed_everything	Seed all random number generators for reproducibility.
current_device	Return the current device.

torchfoo.distributed

is_distributed	Check if distributed mode is initialized
get_world_size	Get number of distributed processes
get_rank	Get rank of the current distributed process
rank_zero_only	Decorator that ensures a function only executes on rank 0.
setup_distributed	Initialize a distributed process group.
cleanup_distributed	Destroy the distributed process group if one is initialized.
parallelize	Decorator that parallelizes a function across multiple devices with distributed setup.
all_reduce_	Perform an in-place differentiable "all reduce" operation
all_reduce_sum_	Perform an in-place differentiable "all reduce sum" operation
all_equal	Check if x is the same on all GPUs
all_concat_jagged	Concatenate a Tensor with varying dim 0 from multiple GPUs
all_concat	Concatenate a Tensor distributed on multiple GPUs.

torchfoo.ddp

make_module_ddp

Wrap a module with DDP and convert BatchNorm to SyncBatchNorm.