Model

class opennmt.models.Model(*args, **kwargs)[source]

Base class for models.

Inherits from: keras.src.engine.base_layer.Layer

Extended by:

property unsupervised

Unsupervised model.

property features_inputter

The inputter producing features.

property labels_inputter

The inputter producing labels.

__repr__()[source]

Returns a description of the model and its submodules.

auto_config(num_replicas=1)[source]

Returns automatic configuration values specific to this model.

Parameters

num_replicas – The number of synchronous model replicas used for the training.

Returns

A partial training configuration.

initialize(data_config, params=None)[source]

Initializes the model from the data configuration.

Parameters

  • data_config – A dictionary containing the data configuration set by the user (e.g. vocabularies, tokenization, pretrained embeddings, etc.).

  • params – A dictionary of hyperparameters.

set_jit_compile(enable)[source]

Allow (or not) this model to use XLA compilation.

build(input_shape)[source]

Creates the variables of the layer (for subclass implementers).

This is a method that implementers of subclasses of Layer or Model can override if they need a state-creation step in-between layer instantiation and layer call. It is invoked automatically before the first execution of call().

This is typically used to create the weights of Layer subclasses (at the discretion of the subclass implementer).

Parameters

input_shape – Instance of TensorShape, or list of instances of TensorShape if the layer expects a list of inputs (one instance per input).

split_features_labels(batch)[source]

Splits a batch from the dataset into features and labels.

__call__(features, labels=None, training=None, step=None)[source]

Runs the model.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

  • training – If True, run in training mode.

  • step – The current training step.

Returns

A tuple containing,

  • The model outputs (usually unscaled probabilities).

  • The model predictions.

abstract call(features, labels=None, training=None, step=None)[source]

Runs the model.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

  • training – If True, run in training mode.

  • step – The current training step.

Returns

A tuple containing,

  • The model outputs (usually unscaled probabilities).

  • The model predictions.

infer(features)[source]

Runs inference on features.

This is a small convenience wrapper around opennmt.models.Model.call().

Parameters

features – A nested structure of features tf.Tensor.

Returns

The model predictions.

evaluate(features, labels)[source]

Evaluates features predictions against labels.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

Returns

A tuple with the loss and the model predictions.

score(features, labels)[source]

Scores labels.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

Returns

The score results.

train(features, labels, optimizer, loss_scale=None)[source]

Computes and applies the gradients for a batch of examples.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

  • optimizer – The optimizer instance (tf.keras.mixed_precision.LossScaleOptimizer is supported).

  • loss_scale – An optional loss scaling factor.

Returns

The loss.

compute_gradients(features, labels, optimizer, loss_scale=None, normalize_loss=True)[source]

Computes the gradients for a batch of examples.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

  • optimizer – The optimizer instance (tf.keras.mixed_precision.LossScaleOptimizer is supported).

  • loss_scale – An optional loss scaling factor.

  • normalize_loss – Normalize the loss by the sample size.

Returns

A tuple containing,

  • The loss.

  • The gradients.

  • The sample size, if normalize_loss is disabled.

compute_training_loss(features, labels, step=None)[source]

Computes the training loss for a batch of examples.

Parameters

  • features – A nested structure of features tf.Tensor.

  • labels – A nested structure of labels tf.Tensor.

  • step – The current training step.

Returns

A tuple containing,

  • The cumulated loss.

  • The sample size (or None if not returned by the model).

abstract compute_loss(outputs, labels, training=True)[source]

Computes the loss.

Parameters

  • outputs – The model outputs (usually unscaled probabilities).

  • labels – The dict of labels tf.Tensor.

  • training – If True, compute the loss for training.

Returns

The loss or a tuple (numerator, train_denominator, stats_denominator) to use a different normalization for training compared to reporting (e.g. batch-normalized for training vs. token-normalized for reporting).

regularize_loss(loss, variables=None)[source]

Regularizes the loss.

Parameters

  • loss – The loss.

  • variables – List of variables.

Returns

The regularized loss.

get_metrics()[source]

Returns the metrics for this model.

Returns

A dictionary of tf.keras.metrics.Metric metrics.

update_metrics(metrics, predictions, labels)[source]

Computes additional metrics on the predictions.

Parameters

  • metrics – A dictionary of metrics to update.

  • predictions – The model predictions.

  • labels – The dict of labels tf.Tensor.

get_optimizer()[source]

Returns the optimizer for this model.

Returns

A tf.keras.optimizers.legacy.Optimizer instance or None if no optimizer is configured.

serve_function()[source]

Returns a function for serving this model.

Returns

A tf.function.

property tflite_mode

Returns True if the model is being traced for TensorFlow Lite.

enable_tflite_mode()[source]

Enable TensorFlow Lite mode for this model.

tflite_function()[source]

Returns the inference function that should be used for TensorFlow Lite.

Returns

A tf.function.

export(export_dir, exporter=None)[source]

Exports the model for serving.

Parameters
create_variables(optimizer=None)[source]

Creates the model variables by running it once.

Parameters

optimizer – If set, also create the optimizer variables.

transfer_weights(new_model, new_optimizer=None, optimizer=None, ignore_weights=None)[source]

Transfers weights (and optionally optimizer slots) from this model to another.

This default implementation assumes that self and new_model have exactly the same variables. Subclasses can override this method to transfer weights to another model type or architecture. For example, opennmt.models.SequenceToSequence can transfer weights to a model with a different vocabulary.

All model and optimizer variables are expected to be initialized.

Parameters

  • new_model – The new model to transfer weights to.

  • new_optimizer – The new optimizer.

  • optimizer – The optimizer used for the current model.

  • ignore_weights – Optional list of weights to not transfer.

map_v1_weights(weights)[source]

Maps current weights to V1 weights.

Parameters

weights – A nested dictionary following the scope names used in V1. The leaves are tuples with the variable value and optionally the optimizer slots.

Returns

A list of tuples associating variables and their V1 equivalent.

export_assets(asset_dir)[source]

Exports additional assets used by this model.

Parameters

asset_dir – The directory where assets can be written.

Returns

A dictionary of additional assets.

visualize(log_dir)[source]

Setups model visualization (e.g. word embedding projections).

Parameters

log_dir – The log directory.

format_prediction(prediction, params=None)[source]

Formats the model prediction for file saving.

Parameters

  • prediction – The model prediction (same structure as the second output of opennmt.models.Model.call()).

  • params – (optional) Dictionary of formatting parameters.

Returns

A string or list of strings.

format_score(score, params=None, stream=None)[source]

Formats the score result for file saving.

Parameters

  • score – The score result (same structure as the output of opennmt.models.Model.score()).

  • params – (optional) Dictionary of formatting parameters.

print_prediction(prediction, params=None, stream=None)[source]

Prints the model prediction.

Parameters

  • prediction – The model prediction (same structure as the second output of opennmt.models.Model.call()).

  • params – (optional) Dictionary of formatting parameters.

  • stream – (optional) The stream to print to.

print_score(score, params=None, stream=None)[source]

Prints the score result.

Parameters

  • score – The score result (same structure as the output of opennmt.models.Model.score()).

  • params – (optional) Dictionary of formatting parameters.

  • stream – (optional) The stream to print to.