Framework

Model

class mammoth.models.NMTModel(encoder, decoder, attention_bridge)

Bases: BaseModel

Core trainable object in OpenNMT. Implements a trainable interface for a simple, generic encoder + decoder model. :param encoder: an encoder object :type encoder: mammoth.modules.layer_stack.StackXcoder :param decoder: a decoder object :type decoder: mammoth.modules.layer_stack.StackXcoder

count_parameters()

Count and log number of parameters in model.

Returns:

• encoder side parameter count

• decoder side parameter count

Return type:

(int, int)

forward(src, decoder_input, src_mask, metadata=None)

Forward propagate a src and tgt pair for training. Possible initialized with a beginning decoder state.

Parameters:

• src (Tensor) – A source sequence passed to encoder. typically for inputs this will be a padded LongTensor of size (len, batch, features). However, may be an image or other generic input depending on encoder.

• tgt (LongTensor) – A target sequence passed to decoder. Size (tgt_len, batch, features).

• lengths (LongTensor) – The src lengths, pre-padding (batch,).

• bptt (Boolean) – A flag indicating if truncated bptt is set. If reset then init_state

Returns:

• decoder output (tgt_len, batch, hidden)

• dictionary attention dists of (tgt_len, batch, src_len)

Return type:

(FloatTensor, dict[str, FloatTensor])

Trainer

class mammoth.Trainer(model, loss_functions, optim, norm_method='sents', accum_count=[1], accum_steps=[0], device_context=None, gpu_verbose_level=0, report_manager=None, model_saver=None, average_decay=0, average_every=1, model_dtype='fp32', earlystopper=None, dropout=[0.3], dropout_steps=[0], task_queue_manager=None, report_stats_from_parameters=False, report_training_accuracy=False, valid_metrics=None, valid_max_length=None, valid_max_batches=None, valid_timeout=None, valid_decode_timeout=None, valid_start=0, vocabs_dict=None, beam_size=1, max_length=100, block_ngram_repeat=0, world_group_sync=None, flops_config=None, report_tflops=True)

Bases: object

Class that controls the training process.

Parameters:

• model (mammoth.models.model.NMTModel) – translation model to train

• loss_functions – ModelDict containing loss functions

• optim (mammoth.utils.optimizers.Optimizer) – the optimizer responsible for update

• data_type (string) – type of the source input: [text]

• norm_method (string) – normalization methods: [sents|tokens]

• accum_count (list) – accumulate gradients this many times.

• accum_steps (list) – steps for accum gradients changes.

• report_manager (mammoth.utils.ReportMgrBase) – the object that creates reports, or None

• model_saver (mammoth.models.ModelSaverBase) – the saver is used to save a checkpoint. Thus nothing will be saved if this parameter is None

train(train_iter, train_steps, save_checkpoint_steps=5000, valid_iter=None, valid_steps=10000, device_context=None)

The main training loop by iterating over train_iter and possibly running validation on valid_iter.

Parameters:

• train_iter – A generator that returns the next training batch.

• train_steps – Run training for this many iterations.

• save_checkpoint_steps – Save a checkpoint every this many iterations.

• valid_iter – A generator that returns the next validation batch.

• valid_steps – Run evaluation every this many iterations.

Returns:

The gathered statistics.

validate(valid_iter, moving_average=None, task=None)

Validate model using autoregressive generation.

Parameters:

• valid_iter – validate data iterator

• moving_average – optional moving average parameters

• task – optional task specification

Returns:

validation loss statistics

Return type:

nmt.Statistics

class mammoth.utils.Statistics(loss=0, n_words=0, n_correct=None)

Bases: object

Accumulator for loss statistics. Currently calculates:

• accuracy

• perplexity

• elapsed time

accumulate_device(loss, n_words, nan_flag, n_src_words=None, metadata=None)

Accumulate one window’s statistics as GPU tensors, WITHOUT syncing.

Everything stays on the GPU; nothing is copied to the CPU here, so this adds no .item() stall to the training loop. Call materialize() to read the running totals back (a single GPU->CPU sync) when they are needed.

Parameters:

• loss – scalar GPU tensor, summed loss for this window.

• n_words – scalar GPU tensor, summed target (non-padding) tokens.

• nan_flag – scalar GPU bool tensor, True if any loss in the window was NaN.

• n_src_words – optional scalar GPU tensor, summed source tokens.

• metadata – optional task metadata; if given, the window loss is also tracked per task (src_lang_tgt_lang).

accuracy()

compute accuracy

static all_gather_stats(stat, max_size=4096)

Gather a Statistics object accross multiple process/nodes

Parameters:

• stat( – obj:Statistics): the statistics object to gather accross all processes/nodes

• max_size (int) – max buffer size to use

Returns:

Statistics, the update stats object

static all_gather_stats_list(stat_list, max_size=4096)

Gather a Statistics list accross all processes/nodes

Parameters:

• stat_list (list([Statistics])) – list of statistics objects to gather accross all processes/nodes

• max_size (int) – max buffer size to use

Returns:

list of updated stats

Return type:

our_stats(list([Statistics]))

elapsed_time()

compute elapsed time

classmethod from_loss_logits_target(loss: float, logits, target, padding_idx)

Alternate constructor for computing the stats from loss, model prediction logits, and target indices. Note that this is heavy. Only use for validation / debug purposes.

log_tensorboard(prefix, writer, learning_rate, patience, step)

display statistics to tensorboard

materialize()

Fold the device-side accumulators into the CPU fields (one sync).

This is the ONLY place the deferred statistics touch the CPU. All pending GPU scalars are stacked into a single tensor and read back with one .tolist(), so a whole reporting interval costs exactly one GPU->CPU synchronisation instead of one per step. No-op (and idempotent) when nothing is pending.

output(step, num_steps, learning_rate, start, metadata=None)

Write out statistics to stdout.

Parameters:

• step (int) – current step

• n_batch (int) – total batches

• start (int) – start time of step.

ppl()

compute perplexity

tflops()

compute TFLOPs/s (teraFLOPs per second)

update(stat, update_n_src_words=False)

Update statistics by suming values with another Statistics object

Parameters:

• stat – another statistic object

• update_n_src_words (bool) – whether to update (sum) n_src_words or not

xent()

compute cross entropy

Optimizer

Optimizers class

class mammoth.utils.optimizers.AdaFactorFairSeq(params, lr=None, eps=(1e-30, 0.001), clip_threshold=1.0, decay_rate=-0.8, beta1=None, weight_decay=0.0, scale_parameter=True, relative_step=True, warmup_init=False)

Bases: Optimizer

Implements Adafactor algorithm.

This implementation is based on: Adafactor: Adaptive Learning Rates with Sublinear Memory Cost (see https://arxiv.org/abs/1804.04235)

Note that this optimizer internally adjusts the learning rate depending on the scale_parameter, relative_step and warmup_init options. To use a manual (external) learning rate schedule you should set scale_parameter=False and relative_step=False.

Parameters:

• params (iterable) – iterable of parameters to optimize or dicts defining parameter groups

• lr (float, optional) – external learning rate (default: None)

• eps (tuple[float, float]) – regularization constans for square gradient and parameter scale respectively (default: (1e-30, 1e-3))

• clip_threshold (float) – threshold of root mean square of final gradient update (default: 1.0)

• decay_rate (float) – coefficient used to compute running averages of square gradient (default: -0.8)

• beta1 (float) – coefficient used for computing running averages of gradient (default: None)

• weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)

• scale_parameter (bool) – if True, learning rate is scaled by root mean square of parameter (default: True)

• relative_step (bool) – if True, time-dependent learning rate is computed instead of external learning rate (default: True)

• warmup_init (bool) – time-dependent learning rate computation depends on whether warm-up initialization is being used (default: False)

step(closure=None)

Performs a single optimization step.

Parameters:

closure (callable, optional) – A closure that reevaluates the model and returns the loss.

class mammoth.utils.optimizers.MultipleOptimizer(suboptimizers, grad_scaler=None)

Bases: object

Separate sub-optimizers for each distributed component. Handles creation of multiple optimizers, grad scaling, restoring from checkpoint, backward, zero_grad, deciding which suboptimizers to step, and reporting.

property amp

True if use torch amp mix precision training.

externally_managed_step(gradient_syncs)

Step through only the trained suboptimizers

load_state_dict(state_dict)

Loads the optimizer from the state dictionary

state_dict()

Returns the state dictionary

property training_step

Global training step, incremented for each minibatch. There may not be any actual parameters that have been trained this number of steps, or any optimizer stepped this number of times.

zero_grad()

Reset the gradient of all sub-optimizers to zero

class mammoth.utils.optimizers.SubOptimizer(optimizer, learning_rate, learning_rate_decay_fn, max_grad_norm=0, grad_scaler=None)

Bases: object

Wraps a base optimizer (torch.optim.Optimizer). Handles learning rate scheduling and grad clipping.

learning_rate()

Returns the current learning rate.

step()

Update the model parameters based on current gradients.

property training_step

The current training step.

zero_grad()

Zero the gradients of optimized parameters.

mammoth.utils.optimizers.exponential_decay(step, rate, decay_steps, start_step=0)

A standard exponential decay, scaling the learning rate by rate every decay_steps steps.

mammoth.utils.optimizers.get_base_optimizer(opts)

Builds the PyTorch optimizer factory.

We use the default parameters for Adam that are suggested by the original paper https://arxiv.org/pdf/1412.6980.pdf These values are also used by other established implementations, e.g. https://www.tensorflow.org/api_docs/python/tf/train/AdamOptimizer https://keras.io/optimizers/ Recently there are slightly different values used in the paper “Attention is all you need” https://arxiv.org/pdf/1706.03762.pdf, particularly the value beta2=0.98 was used there however, beta2=0.999 is still arguably the more established value, so we use that here as well

Parameters:

options. (opts. The dictionary of)

Returns:

A callable that returns torch.optim.Optimizer instances.

mammoth.utils.optimizers.make_learning_rate_decay_fn(opts)

Returns the learning decay function from options.

mammoth.utils.optimizers.noam_decay(step, warmup_steps, model_dim)

Learning rate schedule described in https://arxiv.org/pdf/1706.03762.pdf.

mammoth.utils.optimizers.noamwd_decay(step, warmup_steps, model_dim, rate, decay_steps, start_step=0)

Learning rate schedule optimized for huge batches

mammoth.utils.optimizers.rsqrt_decay(step, warmup_steps)

Decay based on the reciprocal of the step square root.