基于LSTM encoder-decoder模型实现英文转中文的翻译机器

前言

神经网络机器翻译(NMT, neuro machine tranlation)是AIGC发展道路上的一个重要应用。正是对这个应用的研究，发展出了注意力机制，在此基础上产生了AIGC领域的霸主transformer。我们今天先把注意力机制这些东西放一边，介绍一个对机器翻译起到重要里程碑作用的模型：LSTM encoder-decoder模型（sutskever et al. 2014）。根据这篇文章的描述，这个模型不需要特别的优化，就可以取得超过其他NMT模型的效果，所以我们也来动手实现一下，看看是不是真的有这么厉害。

模型

原文作者采用了4层LSTM模型，每层有1000个单元（每个单元有输入门，输出门，遗忘门和细胞状态更新共计4组状态），采用1000维单词向量，纯RNN部分，就有64M参数。同时，在encoder的输出，和decoder的输出后放一个长度为80000的softmax层（因为论文的输出字典长80000），用于softmax的参数量为320M。整个模型共计320M + 64M = 384M。该模型用了8GPU的服务器训练了10天。
模型大概长这样：

按照现在的算力价格，用8张4090的主机训练每小时要花20多块钱，训练一轮下来需要花费小5000，笔者当然没有这么土豪，所以我们会使用一个参数量小得多的模型，主要为了记录整个搭建过程使用到的工具链和技术。另外，由于笔者使用了一个预训练的词向量库，包含了中英文单词共计128万多条，其中中文90多万，英文30多万，要像论文中一样用一个超大的softmax来预测每个词的概率并不现实，因此先使用一个linear层再加上tanh来简化，加快训练过程，只求能看到收敛。使用tanh而不是其他激活性函数的原因是：经过测试，使用tanh收敛效果是最好的。

笔者的模型看起来像这样：

该模型的主要参数如下：
词向量维度：300
LSTM隐藏层个数：900
LSTM层数：4
linear层输入：900
linear层输出：300
模型参数个数如下为：

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
Seq2Seq                                  [1, 11, 300]              --
├─Encoder: 1-1                           [1, 300]                  --
│    └─LSTM: 2-1                         [1, 10, 900]              23,788,800
│    └─Linear: 2-2                       [1, 10, 300]              270,300
│    └─ReLU: 2-3                         [1, 300]                  --
├─Decoder: 1-2                           [1, 11, 300]              --
│    └─LSTM: 2-4                         [1, 11, 900]              23,788,800
│    └─Linear: 2-5                       [1, 11, 300]              270,300
│    └─ReLU: 2-6                         [1, 11, 300]              --
==========================================================================================
Total params: 48,118,200
Trainable params: 48,118,200
Non-trainable params: 0
Total mult-adds (M): 500.11
==========================================================================================
Input size (MB): 0.02
Forward/backward pass size (MB): 0.20
Params size (MB): 192.47
Estimated Total Size (MB): 192.70
==========================================================================================

如果大家希望了解LSTM层的23,788,800个参数如何计算出来，可以参考pytorch源码 pytorch/torch/csrc/api/src/nn/modules/rnn.cpp中方法void RNNImplBase::reset()的实现。笔者如果日后有空也可能会写一写。

3 单词向量及语料

3.1 语料

先说语料，NMT需要大量的平行语料，语料可以从这里获取。另外有个语料天涯网站大全分享给大家。

3.2 词向量

首先需要对句子进行分词，中英文都需要做分词。中文分词工具本例采用jieba，可直接安装。

$ pip install jieba
...
$ python
Python 3.11.6 (tags/v3.11.6:8b6ee5b, Oct  2 2023, 14:57:12) [MSC v.1935 64 bit (AMD64)] on win32
Type