stable_learning_control.utils.mpi_utils.mpi_pytorch

Helper methods for managing Pytorch MPI processes.

Note

This module is not yet used in any of the current algorithms, but is kept here for future reference.

Functions

setup_pytorch_for_mpi()

Avoid slowdowns caused by each separate process's PyTorch using more than its

mpi_avg_grads(module)

Average contents of gradient buffers across MPI processes.

sync_params(module)

Sync all parameters of module across all MPI processes.

Module Contents

stable_learning_control.utils.mpi_utils.mpi_pytorch.setup_pytorch_for_mpi()[source]

Avoid slowdowns caused by each separate process’s PyTorch using more than its fair share of CPU resources.

stable_learning_control.utils.mpi_utils.mpi_pytorch.mpi_avg_grads(module)[source]

Average contents of gradient buffers across MPI processes.

Parameters:

module (object) – Python object for which you want to average the gradients.

stable_learning_control.utils.mpi_utils.mpi_pytorch.sync_params(module)[source]

Sync all parameters of module across all MPI processes.

Parameters:

module (object) – Python object for which you want to average the gradients.