mlbench_core.utils¶
pytorch¶
FCGraph¶
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
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]
Inference¶
Translator¶
BeamSearch¶
- class mlbench_core.utils.pytorch.inference.beam_search.SequenceGenerator(model, beam_size=5, max_seq_len=100, len_norm_factor=0.6, len_norm_const=5, cov_penalty_factor=0.1)[source]¶
Generator for the autoregressive inference with beam search decoding.
Beam search decoding supports coverage penalty and length normalization. For details, refer to Section 7 of the GNMT paper (https://arxiv.org/pdf/1609.08144.pdf).
- Parameters
model – model which implements generate method
beam_size (int) – decoder beam size
max_seq_len (int) – maximum decoder sequence length
len_norm_factor (float) – length normalization factor
len_norm_const (float) – length normalization constant
cov_penalty_factor (float) – coverage penalty factor
- beam_search(self, batch_size, initial_input, initial_context=None)[source]¶
Beam Search decoder.
- Parameters
batch_size (int) – decoder batch size
initial_input (torch.tensor) – initial input, usually tensor of BOS tokens
initial_context (torch.tensor) – initial context, usually [encoder_context, src_seq_lengths, None]
- Returns: (translation, lengths, counter)
translation: (batch_size, max_seq_len) - indices of target tokens lengths: (batch_size) - lengths of generated translations counter: number of iterations of the decoding loop
- greedy_search(self, batch_size, initial_input, initial_context=None)[source]¶
Greedy decoder.
- Parameters
batch_size (int) – decoder batch size
initial_input (torch.tensor) – initial input, usually tensor of BOS tokens
initial_context (torch.tensor) – initial context, usually [encoder_context, src_seq_lengths, None]
- Returns: (translation, lengths, counter)
translation: (batch_size, max_seq_len) - indices of target tokens lengths: (batch_size) - lengths of generated translations counter: number of iterations of the decoding loop
tensorflow¶
Initialize environment for pytorch.