PyTorch 中的序列到序列模型

这是一个用于在 PyTorch 中训练序列到序列模型的完整工具包。它包含多个模型以及用于训练和推理的代码。

使用此代码，您可以训练：

• 神经机器翻译 (NMT) 模型
• 语言模型
• 图像描述生成
• Skip-Thought 句子表示
• 以及其他更多……

安装

git clone --recursive https://github.com/eladhoffer/seq2seq.pytorch
cd seq2seq.pytorch; python setup.py develop

模型

当前可用的模型包括：

• 简单的循环序列到序列模型
• 带注意力机制解码器的循环序列到序列模型
• Google 神经机器翻译（GNMT）循环模型
• Transformer——一种完全基于注意力机制的模型，源自论文《Attention Is All You Need》（https://arxiv.org/abs/1706.03762）

数据集

当前可用的数据集有：

• WMT16
• WMT17
• OpenSubtitles 2016
• COCO 图像标注
• Conceptual Captions

所有数据集均可使用三种分词方法进行标记化：

• 基于字符的分词
• 基于单词的分词
• 字节对编码（BPE），由 bpe 提出，可选择不同的词汇表大小。

选择分词方法后，将生成并保存词汇表，以供后续推理使用。

训练方法

这些模型可以通过多种方式进行训练：

• 基本序列到序列模型：给定编码后的输入序列，生成（解码）输出序列。训练采用教师强制法。
• 多任务序列到序列模型：同时训练多个任务（例如多种语言），将数据序列既作为编码器的输入，又作为解码器的输出。
• 图像到序列模型：用于训练图像描述生成器。

使用

示例训练脚本位于 scripts 文件夹中。推理示例则位于 examples 文件夹中。

• 以下示例按照原始论文的设置，在 WMT16 数据集上训练一个 Transformer 模型：

DATASET=${1:-"WMT16_de_en"}
DATASET_DIR=${2:-"./data/wmt16_de_en"}
OUTPUT_DIR=${3:-"./results"}

WARMUP="4000"
LR0="512**(-0.5)"

python main.py \
  --save transformer \
  --dataset ${DATASET} \
  --dataset-dir ${DATASET_DIR} \
  --results-dir ${OUTPUT_DIR} \
  --model Transformer \
  --model-config "{'num_layers': 6, 'hidden_size': 512, 'num_heads': 8, 'inner_linear': 2048}" \
  --data-config "{'moses_pretok': True, 'tokenization':'bpe', 'num_symbols':32000, 'shared_vocab':True}" \
  --b 128 \
  --max-length 100 \
  --device-ids 0 \
  --label-smoothing 0.1 \
  --trainer Seq2SeqTrainer \
  --optimization-config "[{'step_lambda':
                          \"lambda t: { \
                              'optimizer': 'Adam', \
                              'lr': ${LR0} * min(t ** -0.5, t * ${WARMUP} ** -1.5), \
                              'betas': (0.9, 0.98), 'eps':1e-9}\"
                          }]"

• 以下示例训练一个基于注意力机制的 LSTM 模型，编码器和解码器各包含 3 层：

python main.py \
  --save de_en_wmt17 \
  --dataset ${DATASET} \
  --dataset-dir ${DATASET_DIR} \
  --results-dir ${OUTPUT_DIR} \
  --model RecurrentAttentionSeq2Seq \
  --model-config "{'hidden_size': 512, 'dropout': 0.2, \
                   'tie_embedding': True, 'transfer_hidden': False, \
                   'encoder': {'num_layers': 3, 'bidirectional': True, 'num_bidirectional': 1, 'context_transform': 512}, \
                   'decoder': {'num_layers': 3, 'concat_attention': True,\
                               'attention': {'mode': 'dot_prod', 'dropout': 0, 'output_transform': True, 'output_nonlinearity': 'relu'}}}" \
  --data-config "{'moses_pretok': True, 'tokenization':'bpe', 'num_symbols':32000, 'shared_vocab':True}" \
  --b 128 \
  --max-length 80 \
  --device-ids 0 \
  --trainer Seq2SeqTrainer \
  --optimization-config "[{'epoch': 0, 'optimizer': 'Adam', 'lr': 1e-3},
                          {'epoch': 6, 'lr': 5e-4},
                          {'epoch': 8, 'lr':1e-4},
                          {'epoch': 10, 'lr': 5e-5},
                          {'epoch': 12, 'lr': 1e-5}]" \

seq2seq.pytorch 快速上手指南

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux 或 macOS（Windows 需通过 WSL 运行）
• Python 版本：推荐 Python 3.6+
• 核心依赖：
  • PyTorch (最新稳定版)
  • Git
• 国内加速建议：
  • 安装 PyTorch 时，推荐使用清华或中科大镜像源。例如：

    pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple
    

  • 克隆代码库时，若 GitHub 连接缓慢，可使用 Gitee 镜像（如有）或配置代理。

安装步骤

本项目支持递归克隆以包含所有子模块，并通过 setup.py 进行开发模式安装。

1. 克隆仓库

   git clone --recursive https://github.com/eladhoffer/seq2seq.pytorch
   

   提示：如果克隆速度较慢，可尝试使用国内镜像源（如 https://gitee.com/mirrors/seq2seq.pytorch.git，若可用）或设置 GIT_SSL_NO_VERIFY=true。

2. 进入目录并安装

   cd seq2seq.pytorch
   python setup.py develop
   

   此命令会将包链接到当前环境，方便后续修改代码并立即生效。

基本使用

安装完成后，您可以使用 main.py 脚本进行模型训练。以下是一个最基础的示例，展示如何训练一个 Transformer 模型（基于 WMT16 德语 - 英语数据集）。

训练示例

将以下命令保存为脚本或直接终端运行。该示例使用了 BPE 分词、共享词表以及论文中推荐的优化器策略。

DATASET=${1:-"WMT16_de_en"}
DATASET_DIR=${2:-"./data/wmt16_de_en"}
OUTPUT_DIR=${3:-"./results"}

WARMUP="4000"
LR0="512**(-0.5)"

python main.py \
  --save transformer \
  --dataset ${DATASET} \
  --dataset-dir ${DATASET_DIR} \
  --results-dir ${OUTPUT_DIR} \
  --model Transformer \
  --model-config "{'num_layers': 6, 'hidden_size': 512, 'num_heads': 8, 'inner_linear': 2048}" \
  --data-config "{'moses_pretok': True, 'tokenization':'bpe', 'num_symbols':32000, 'shared_vocab':True}" \
  --b 128 \
  --max-length 100 \
  --device-ids 0 \
  --label-smoothing 0.1 \
  --trainer Seq2SeqTrainer \
  --optimization-config "[{'step_lambda':
                          \"lambda t: { \
                              'optimizer': 'Adam', \
                              'lr': ${LR0} * min(t ** -0.5, t * ${WARMUP} ** -1.5), \
                              'betas': (0.9, 0.98), 'eps':1e-9}\"
                          }]"

关键参数说明

• --model: 选择模型架构，可选 Transformer, RecurrentAttentionSeq2Seq, SimpleSeq2Seq 等。
• --tokenization: 数据预处理方式，支持 bpe (Byte-pair-encoding), word, char。
• --device-ids: 指定使用的 GPU ID，多卡训练可设为 0,1,2,3。
• scripts 文件夹：包含更多预定义的训练脚本。
• examples 文件夹：包含推理（Inference）示例代码。

完成训练后，模型权重将保存在 --results-dir 指定的目录中，可用于后续的推理任务。