# translator

**Repository Path**: chen-junfa/translator

## Basic Information

- **Project Name**: translator
- **Description**: 纯torch的transformer机器翻译入门项目
- **Primary Language**: Python
- **License**: MIT
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 2
- **Forks**: 1
- **Created**: 2025-08-09
- **Last Updated**: 2026-01-07

## Categories & Tags

**Categories**: Uncategorized

**Tags**: Deep-learning, Transformer, torch, machine-translation, 入门项目

## README

# 基于Transformer的机器翻译（学习版）

## 介绍

本项目是一个基于transformer的机器翻译（中译英）练手项目，使用纯torch完成，未依赖高度封装的成品库，所用库皆为最新（截止2025-08-10），并且数据已经准备好，可作为transformer的入门示例。

## 项目结构

```python
.
├── checkpoints
│   └── translator_best.pt      # 模型权重文件
├── data
│   ├── data_xxxx.jsonl         # 数据集文件
│   ├── dataset.py              # dataloder设计
│   └── __init__.py
├── model
│   ├── __init__.py
│   ├── model.py                # 模型主体结构设计
│   └── modules.py              # 模型的各模块设计
├── README.md                   # 项目文档
├── requirements.txt            # python依赖
├── results.txt                 # 测试结果
├── test.py                     # 测试脚本
├── tokenizer                   # 分词器设计（简单字符级分词器）
│   ├── build_vocab.py          # 从数据集中统计字符频率，构建词表
│   ├── function.py             # 分词器具体功能实现
│   ├── __init__.py
│   └── vocab.json              # 词表文件
└── train.py                    # 训练脚本
```

## 快速开始

1. 安装conda(一种python环境管理工具)，具体介绍、下载和用法搜索关键词minicoda(更轻量)或者anaconda(更齐全)，本项目建议miniconda即可。
2. 创建python环境 `conda create -n translator python=3.12`并启用它 `conda activate translator`
3. 克隆该仓库并进入仓库根目录
4. 安装依赖 `pip install -r requirements.txt`
5. 训练 `python train.py`（注：默认配置下需要至少需要配备12G显存的nvidiaGPU，训练10个epoch在RTX 3060 12G上耗时约 30h）（显存不足降低batch_size的值即可）
6. 测试 `python test.py`（默认配置下在RTX 3060 12G上测试耗时约 15min）

仓库中自带一份训练好的权重文件以及测试结果

## 数据准备

本项目使用[WMT中英机器翻译训练集](https://www.modelscope.cn/datasets/iic/WMT-Chinese-to-English-Machine-Translation-Training-Corpus/files)来训练模型。完整数据集过于庞大，本项目只截取了其中2,500,000条样本，保存在 `data/data_xxxx.jsonl`中。

## tokenizer设计

`tokenizer负责将自然语句转化为id列表（句子=>片段的列表=>id的列表），或者反之将id列表转化为句子。

本项目采用最简单的实现方法——将句子按字符拆分，同时为每个字符分配一个数字作为id，从而转化为id列表。

构建词表的代码见 `tokenizer/build_vocab.py`。

tokenizer的功能实现见 `tokenizer/function.py`。

## 模型设计

### 整体架构

整体使用encoder-decoder架构，即将输入进行向量化，然后通过与已经输出的向量进行self_attention和与输入向量进行cross_attention来计算下一个输出token。

### 嵌入层

经过 `tokenizer.encode`之后，语句被处理为了id序列，但是此时的id序列的元素仅仅是整数，携带的信息量太少，模型很难理解和区分他们，所以我们需要为每个token分配一个信息量更大的数据（比如一个向量，称之为embedding）。此时，id相当于token的索引，embedding相当于token的详细介绍。具体操作方法为设置一个V*D的张量E（V为词表大小，D为嵌入维度，即用多少个数据来描述一个token），`E[token_id]`即为 `token_embedding`。pytorch中已经将该功能封装，使用 `torch.nn.Embedding`即可，具体用法见代码 `model/modules.py`

### SelfAttention

Seq2Seq模型的关键机制，最简示例具体逻辑见 `model/modules.py`的 `SelfAtten_Simple`类，数学上看是将每个输入向量x通过三个不同的矩阵映射为q/k/v三个向量，每个x用自己的q与其他人的k做内积，得到一个类似“匹配程度”的数，然后以这个匹配程度为权重加权平均其他人的v，从而获得自己对应的输出向量。该算法无需固定序列长度，并且顺序无关，同时每个token的输出向量既包含自己的信息，又包含全局信息。

进一步拓展为多头注意力机制，相当于同时有多个不同的简单self_atten，输出多个结果，拼接再经过线性层融合后得到最终输出。一般我们会通过编程技巧将多个self_atten压成一个模型，从而提高并行度，充分利用GPU性能。同时为了匹配我们实际使用时逐个输出的模式，我们不能让一个token去包含它后面的token的v，相当于需要手动将这部分的分数置零，于是需要引入mask。引入mask的多头注意力机制具体实现见 `model/modules.py`的 `MultiHeadAtten`类。

### EncodingLayer

上述的attention模块只包括线性部分，无法描述非线性关系，需要结合非线性层提高拟合能力。参考transformer论文的设计，结合注意力，前馈网络，layernorm和残差连接，封装一个上层类 `EncodingLayer`作为 `Encoder`的基本组件。

### 位置编码

attention算法数学上是位置无关的，但是在我们的翻译任务中，token的顺序是有十分重要的意义的，所以我们需要给token额外添加位置信息，即进行位置编码。本项目中我们使用Transformer论文中的正弦位置编码。

$$
\begin{align}
\text{PE}_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}}\right) \\
\text{PE}_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{\frac{2i}{d_{\text{model}}}}}\right)
\end{align}
$$

其核心是利用傅里叶变换的的强大表示潜力，让模型能够自己探索出合适的使用方法来应对任务。

具体函数实现见 `model/modules.py`的 `PositionalEncoding`类。

### Encoder

完成上述组件即可搭建 `Encoder`，`Embedding+PositionalEncoding=>N层EncodingLayer`即可。(逻辑上如此，但是实际代码中本项目使用共享embedding层，所以Encoder内部不带embedding)

### CrossAttention

在原有多头注意力的基础上，额外使用自己的q去匹配另一个序列内的tokens的kv，即为交叉注意力，其目的是为了让一个序列去融合另一个序列的信息。具体函数实现见 `model/modules.py`的 `CrossAtten`类。

### Decoder

类似 `Encoder`，组合完 `DecodingLayer`之后即可组合 `Decoder`，注意mask构建即可。具体函数实现见 `model/modules.py`的 `Decoder`类。

注：本项目的Encoder和Decoder为了简洁均使用半屏蔽mask，数学上实现了屏蔽效果但是会使矩阵更加稠密，略微降低前向传播速度。

### 最终模型结构

见 `model/model.py`。

## Dataset

具体见 `data/dataset.py`，本项目使用流式读取节省内存占用。同时拆分了训练集、验证集和测试集，这是深度学习中的推荐做法。

一般我们会使用训练集的数据来训练模型（参与反向传播），每训练一定epoch数或者步数就会保存一次模型权重（也可以额外保存其他信息），保存下来的这些权重称之为“检查点”（checkpoints），从中选取在验证集上loss最小的一个版本最为最终版本，然后用其在测试集上的得分（可以说loss也可以是其他评估性能的指标）作为模型质的最终评价。

| 选验证集loss最低的版本 | 推荐做法                                                                                                         |
| ---------------------- | ---------------------------------------------------------------------------------------------------------------- |
| 选训练集loss最低的版本 | 有过拟合风险，或者说泛化性可能不好<br />（模型可能只是背下训练集的答案和套路，应对训练集以外的任务就会原形毕露） |
| 选测试集loss最低的版本 | 指标有水分，类似得到“最高分”而非“平均分”                                                                     |

|              | 验证集loss大 | 验证集loss小 |
| ------------ | ------------ | ------------ |
| 训练集loss大 | 欠拟合       | 狗运         |
| 训练集loss小 | 过拟合       | 理想结果     |


## 训练
见`train.py`，使用了早停（early stop）来防止过拟合，并且启用了混合精度训练来加速。
```python
with torch.amp.autocast(device_type=device.type): #torch自动处理混合精度的上下文管理器
    ......
```

## 测试
见`test.py`，使用`nltk`库来计算BELU分数，最终BELU得分0.3885，并且输出显示哪怕使用的是最简陋的字符集分词器，模型在经过大量训练之后仍然能够理解到基本的单词、语法和中英对应关系，翻译效果勉强可看。