OpenNMT-tf

OpenNMT-tf 是一款基于 TensorFlow 2 的通用序列学习工具包。虽然神经机器翻译是其主要目标任务，但该工具包也旨在更广泛地支持以下多种任务：

• 序列到序列映射
• 序列标注
• 序列分类
• 语言建模

该项目以生产级为目标，具备【向后兼容性保障】（详见：https://github.com/OpenNMT/OpenNMT-tf/blob/master/CHANGELOG.md）。

核心功能

模块化模型架构

模型的定义采用代码形式，便于用户自定义模型架构并覆盖默认行为。例如，以下示例定义了一个序列到序列模型，该模型包含 2 个串联的输入特征、一个自注意力编码器，以及一个共享输入和输出嵌入的注意力 RNN 解码器：

opennmt.models.SequenceToSequence(
    source_inputter=opennmt.inputters.ParallelInputter(
        [
            opennmt.inputters.WordEmbedder(embedding_size=256),
            opennmt.inputters.WordEmbedder(embedding_size=256),
        ],
        reducer=opennmt.layers.ConcatReducer(axis=-1),
    ),
    target_inputter=opennmt.inputters.WordEmbedder(embedding_size=512),
    encoder=opennmt.encoders.SelfAttentionEncoder(num_layers=6),
    decoder=opennmt.decoders.AttentionalRNNDecoder(
        num_layers=4,
        num_units=512,
        attention_mechanism_class=tfa.seq2seq.LuongAttention,
    ),
    share_embeddings=opennmt.models.EmbeddingsSharingLevel.TARGET,
)

opennmt 包还提供了其他构建模块，可用于设计各类模型：

• 多个输入特征（如：ParallelInputter）
• 混合嵌入表示（如：MixedInputter）
• 多源上下文（如：ParallelInputter）
• 串行编码器（如：SequentialEncoder）或多列编码器（如：ParallelEncoder）
• 混合序列到序列模型（如：SequenceToSequence）

标准模型，例如 Transformer，已在【模型目录】中进行定义，并且无需额外配置即可直接使用。

• 更多关于模型配置的信息，请参阅【文档】（https://opennmt.net/OpenNMT-tf/package/overview.html）。

全面集成 TensorFlow 2

OpenNMT-tf 与 TensorFlow 2 生态系统实现了深度整合：

• 可重用层，可扩展自 tf.keras.layers.Layer（如：Layer）
• 支持多 GPU 训练，可通过 tf.distribute 实现分布式训练，同时借助 Horovod 进行分布式训练（GitHub: https://github.com/horovod/horovod）
• 支持混合精度训练，通过 tf.keras.mixed_precision 实现
• 提供可视化功能，如 TensorBoard（GitHub: https://www.tensorflow.org/tensorboard）
• 使用 tf.function 进行图追踪，可导出为 SavedModel（GitHub: https://opennmt.net/OpenNMT-tf/serving.html），并通过 TensorFlow Serving（GitHub: https://github.com/OpenNMT/OpenNMT-tf/tree/master/examples/serving/tensorflow_serving）或 Python 方式进行部署。

与 CTranslate2 的兼容性

[CTranslate2] 是一款专为 OpenNMT 模型优化的推理引擎，具备高效的 CPU 和 GPU 执行能力、模型量化、并行翻译、动态内存管理、交互式解码等功能！OpenNMT-tf 能够自动将模型导出为 CTranslate2 所需的格式。

动态数据管道

OpenNMT-tf 不需要在训练前对数据进行编译，而是可以直接读取文本文件，并在训练过程中根据需要对数据进行预处理。这使得我们能够实现【即时分词】（如：tokenization）以及通过注入随机噪声来实现数据增强。

模型微调

OpenNMT-tf 支持模型微调流程：

• 模型权重可以迁移到新的词汇表中，例如在针对域内数据进行微调之前，注入领域术语
• 通过【对比学习】（如：https://ai.google/research/pubs/pub48253）来减少单词遗漏错误

源目标对齐

序列到序列模型可通过【引导对齐】（如：arXiv:1607.01628）进行训练，而对齐信息会作为翻译 API 的一部分返回。

---

OpenNMT-tf 同时实现了大多数用于训练和评估序列模型的常用技术，例如：

• 在训练过程中自动进行评估

• 多种解码策略：贪婪搜索、束搜索、随机采样

• N 个最佳结果的重新加权排序

• 梯度累积

• 定期采样

• 检查点平均值

• ……以及其他更多功能！

• 请参阅【文档】（https://opennmt.net/OpenNMT-tf/）以了解如何使用这些功能。

使用方法

OpenNMT-tf 需要：

• Python 3.7 或更高版本
• TensorFlow 2.6、2.7、2.8、2.9、2.10、2.11、2.12 或 2.13

我们建议使用 pip 进行安装：

pip install --upgrade pip
pip install OpenNMT-tf

• 请参阅【文档】（https://opennmt.net/OpenNMT-tf/installation.html）以获取更多信息。*

命令行

OpenNMT-tf 提供了多个命令行工具，用于准备数据、训练和评估模型。

对于所有涉及模型运行的任务，OpenNMT-tf 都使用一个独特的入口点：onmt-main。典型的 OpenNMT-tf 运行由三个部分组成：

• 模型 类型
• 参数，以 YAML 文件的形式描述
• 运行 类型，例如 train、eval、infer、export、score、average_checkpoints 或 update_vocab

这些参数会被传递给主脚本：

onmt-main --model_type <model> --config <config_file.yml> --auto_config <run_type> <run_options>

• 如需了解更多关于 OpenNMT-tf 的使用方法及示例，请访问【我们的文档】（https://opennmt.net/OpenNMT-tf）。

图书馆

OpenNMT-tf 还提供了【清晰定义且稳定的 API】（https://opennmt.net/OpenNMT-tf/package/overview.html），从高级训练工具到低级模型层，以及数据集转换，一应俱全。

例如，只需寥寥数行代码，即可使用 Runner 类来训练和评估模型：

import opennmt

config = {
    "model_dir": "/data/wmt-ende/checkpoints/",
    "data": {
        "source_vocabulary": "/data/wmt-ende/joint-vocab.txt",
        "target_vocabulary": "/data/wmt-ende/joint-vocab.txt",
        "train_features_file": "/data/wmt-ende/train.en",
        "train_labels_file": "/data/wmt-ende/train.de",
        "eval_features_file": "/data/wmt-ende/valid.en",
        "eval_labels_file": "/data/wmt-ende/valid.de",
    }
}

model = opennmt.models.TransformerBase()
runner = opennmt.Runner(model, config, auto_config=True)
runner.train(num_devices=2, with_eval=True)

以下是另一个示例，利用 OpenNMT-tf 以高效的束搜索方式运行自注意力解码器：

decoder = opennmt.decoders.SelfAttentionDecoder(num_layers=6, vocab_size=32000)

initial_state = decoder.initial_state(
    memory=memory, memory_sequence_length=memory_sequence_length
)

batch_size = tf.shape(memory)[0]
start_ids = tf.fill([batch_size], opennmt.START_OF_SENTENCE_ID)

decoding_result = decoder.dynamic_decode(
    target_embedding,
    start_ids=start_ids,
    initial_state=initial_state,
    decoding_strategy=opennmt.utils.BeamSearch(4),
)

更多使用 OpenNMT-tf 作为库的示例，可在线查阅：

• examples/library 目录中，收录了更多以 OpenNMT-tf 为库的示例。
• nmt-wizard-docker 使用高级的 opennmt.Runner API，将 OpenNMT-tf 与自定义接口相结合，用于训练、翻译及服务。

*如需全面了解各类 API，请参阅【软件包文档】（https://opennmt.net/OpenNMT-tf/package/overview.html）。

其他资源

• 【文档】（https://opennmt.net/OpenNMT-tf）
• 【论坛】（https://forum.opennmt.net）
• 【Gitter】（https://gitter.im/OpenNMT/OpenNMT-tf）

OpenNMT-tf 快速上手指南

环境准备

• 操作系统：Linux / macOS / Windows
• Python：3.7 及以上
• TensorFlow：2.6 ~ 2.13 任一版本
• 硬件：CPU 即可运行；GPU 训练需 CUDA 11.2+（推荐 NVIDIA 驱动 ≥ 470）

安装步骤

# 1. 升级 pip
pip install --upgrade pip

# 2. 安装 OpenNMT-tf（国内可换清华源加速）
pip install OpenNMT-tf
# 如需 GPU 支持，先装对应版本 TensorFlow
# pip install tensorflow==2.12

基本使用

1. 准备数据（示例：英德翻译）

train.en   # 英文源文件
train.de   # 德文目标文件
valid.en   # 验证源文件
valid.de   # 验证目标文件

2. 创建配置文件 config.yml

model_dir: ./run/

data:
  source_vocabulary: vocab.src
  target_vocabulary: vocab.tgt
  train_features_file: train.en
  train_labels_file: train.de
  eval_features_file: valid.en
  eval_labels_file: valid.de

3. 一键训练

# 使用内置 Transformer 模型
onmt-main --model_type TransformerBase \
          --config config.yml \
          --auto_config train --with_eval

4. 推理翻译

echo "Hello world" > input.txt
onmt-main --model_type TransformerBase \
          --config config.yml \
          --auto_config infer --features_file input.txt

输出结果将保存在 run/predictions.txt。