mlbench_core.utils

pytorch

FCGraph

class mlbench_core.utils.pytorch.FCGraph(rank, world_size, use_cuda=False)[source]

Fully-Connected Network Graph

Parameters

config (dict) – a global object containing all of the config.

initialize_backends

mlbench_core.utils.pytorch.initialize_backends(comm_backend='mpi', hosts=None, rank=- 1, logging_level='INFO', logging_file='/mlbench.log', use_cuda=False, seed=None, cudnn_deterministic=False, ckpt_run_dir='/checkpoints', delete_existing_ckpts=False)[source]

Initializes the backends.

Sets up logging, sets up pytorch and configures paths correctly.

Parameters

config (types.SimpleNamespace) – a global object containing all of the config.

Returns

a global object containing all of the config.

Return type

(types.SimpleNamespace)

Checkpointer

class mlbench_core.utils.pytorch.checkpoint.Checkpointer(ckpt_run_dir, rank, freq=CheckpointFreq.BEST, save_stats=True)[source]

A class for handling checkpoint saving and loading.

Parameters
  • ckpt_run_dir (str) – The path of the checkpoint directory.

  • rank (int) – The rank of the eurrent worker.

  • freq (int) – The frequency of checkpointing. Default: CheckpointFreq.BEST

  • save_stats (bool) – Save stats to additional text files. Default: True

helpers

Helper functions.

mlbench_core.utils.pytorch.helpers.config_logging(logging_level='INFO', logging_file='/mlbench.log')[source]

Setup logging modules. A stream handler and file handler are added to default logger mlbench.

Parameters
  • logging_level (str) – Log level

  • logging_file (str) – Log file

mlbench_core.utils.pytorch.helpers.config_pytorch(use_cuda=False, seed=None, cudnn_deterministic=False)[source]

Config pytorch packages.

Fix random number for packages and initialize distributed environment for pytorch. Setup cuda environment for pytorch.

Parameters
  • use_cuda (bool) – Use CUDA acceleration

  • seed (int | None) – Random seed to use

  • cudnn_deterministic (bool) – Set cudnn.determenistic=True

Returns

FCGraph): The rank, world size, and network graph

Return type

(int, int, obj

mlbench_core.utils.pytorch.helpers.config_path(ckpt_run_dir, delete_existing_ckpts=False)[source]

Config the path used during the experiments.

utils

mlbench_core.utils.pytorch.utils.pack_tensors(tensors, use_cuda=False)[source]

Packs a list of tensors into one 1-dimensional tensor.

Parameters
  • tensors (list[torch.Tensor]) – The tensors to pack

  • use_cuda (bool) – Whether the resulting tensor should be on cuda

Returns

The flattened tensors, the list start indices of each packed tensor, and the original shape of each tensor.

Those values are used to then unpack the tensor

Return type

(torch.Tensor, list[int], list[(int, int)])

mlbench_core.utils.pytorch.utils.unpack_tensors(aggregated, indices, sizes)[source]

Unpacks a 1-dimensional tensor into a list of tensors

Parameters
  • aggregated (torch.Tensor) – The 1-dimensional tensor

  • indices (List[Int]) – The start index of each tensor

  • sizes (List[(Int, Int)]) – The size of each resulting tensor

Returns

The unpacked tensors

Return type

List[torch.Tensor]

mlbench_core.utils.pytorch.utils.orthogonalize(matrix, eps=torch.FloatTensor([1e-16]))[source]

Function used to orthogonalize a matrix.

Parameters
  • matrix (torch.Tensor) – Matrix to orthogonalize

  • eps (torch.FloatTensor) – Used to avoid division by zero (default: 1e-16)

Inference

Translator

tensorflow

Initialize environment for pytorch.