opennmt.models.model module

Base class for models.

class opennmt.models.model.Model(name, features_inputter=None, labels_inputter=None, daisy_chain_variables=False, dtype=None, examples_inputter=None)[source]

Bases: tensorflow.python.keras.engine.base_layer.Layer

Base class for models.

unsupervised

Unsupervised model.

features_inputter

The inputter producing features.

labels_inputter

The inputter producing labels.

auto_config(num_devices=1)[source]

Returns automatic configuration values specific to this model.

Parameters:num_devices – The number of devices used for the training.
Returns:A partial training configuration.
initialize(metadata)[source]

Initializes the model from the data configuration.

Parameters:metadata – A dictionary containing additional data configuration set by the user (e.g. vocabularies, tokenization, pretrained embeddings, etc.).
__call__(features, labels, params, mode, config=None)[source]

Calls the model function.

Returns:
The model outputs (usually unscaled probabilities).
Optional if mode is tf.estimator.ModeKeys.PREDICT.
predictions: The model predictions.
Optional if mode is tf.estimator.ModeKeys.TRAIN.
Return type:outputs

See also

tf.estimator.Estimator ‘s model_fn argument for more details about the arguments of this function.

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

Computes the loss.

Parameters:
  • outputs – The model outputs (usually unscaled probabilities).
  • labels – The dict of labels tf.Tensor.
  • training – Compute training loss.
  • params – A dictionary of hyperparameters.
Returns:

The loss or a tuple containing the computed loss and the loss to display.

optimize_loss(loss, params=None, hvd=None)[source]

Returns the loss optimization op.

Parameters:
  • loss – The loss to optimize.
  • params – A dictionary of hyperparameters.
  • hvd – Optional Horovod object.
Returns:

The training op and optionally a list of extra variables to initialize.

compute_metrics(predictions, labels)[source]

Computes additional metrics on the predictions.

Parameters:
  • predictions – The model predictions.
  • labels – The dict of labels tf.Tensor.
Returns:

A dict of metrics. See the eval_metric_ops field of tf.estimator.EstimatorSpec.

input_fn(mode, batch_size, metadata=None, features_file=None, labels_file=None, batch_type='examples', batch_multiplier=1, bucket_width=None, single_pass=False, num_threads=None, sample_buffer_size=None, prefetch_buffer_size=None, maximum_features_length=None, maximum_labels_length=None, num_shards=1, shard_index=0)[source]

Returns an input function.

Parameters:
  • mode – A tf.estimator.ModeKeys mode.
  • batch_size – The batch size to use.
  • metadata – A dictionary containing additional metadata set by the user. Required if Model.initialize() has not been called.
  • features_file – The file containing input features.
  • labels_file – The file containing output labels.
  • batch_type – The training batching stragety to use: can be “examples” or “tokens”.
  • batch_multiplier – The batch size multiplier to prepare splitting accross replicated graph parts.
  • bucket_width – The width of the length buckets to select batch candidates from. None to not constrain batch formation.
  • single_pass – If True, makes a single pass over the training data.
  • num_threads – The number of elements processed in parallel.
  • sample_buffer_size – The number of elements from which to sample.
  • prefetch_buffer_size – The number of batches to prefetch asynchronously. If None, use an automatically tuned value on TensorFlow 1.8+ and 1 on older versions.
  • maximum_features_length – The maximum length or list of maximum lengths of the features sequence(s). None to not constrain the length.
  • maximum_labels_length – The maximum length of the labels sequence. None to not constrain the length.
  • num_shards – The number of data shards (usually the number of workers in a distributed setting).
  • shard_index – The shard index this input pipeline should read from.
Returns:

A callable that returns the next element.

See also

tf.estimator.Estimator.

model_fn(num_devices=1, eval_prediction_hooks_fn=None, devices=None, hvd=None)[source]

Returns the model function.

Parameters:
  • num_devices – The number of devices used for training.
  • eval_prediction_hooks_fn – A callable that takes the model predictions during evaluation and return an iterable of evaluation hooks (e.g. for saving predictions on disk, running external evaluators, etc.).
  • devices – The list of devices used for training, if known.
  • hvd – Optional Horovod object.

See also

tf.estimator.Estimator ‘s model_fn argument for more details about arguments and the returned value.

serving_input_fn(metadata=None)[source]

Returns the serving input function.

Parameters:metadata – A dictionary containing additional metadata set by the user. Required if Model.initialize() has not been called.
Returns:A callable that returns a tf.estimator.export.ServingInputReceiver.
get_assets(metadata=None, asset_dir=None)[source]

Returns additional assets used by this model.

Parameters:
  • metadata – A dictionary containing additional metadata set by the user. Required if Model.initialize() has not been called.
  • asset_dir – The directory where assets can be written.
Returns:

A dictionary of additional assets.

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

Prints the model prediction.

Parameters:
  • prediction – The evaluated prediction.
  • params – (optional) Dictionary of formatting parameters.
  • stream – (optional) The stream to print to.