【AI 进阶笔记】Transformer 中英翻译实现

原创

繁依Fanyi

发布于 2025-03-27 12:17:16

20800

代码可运行

运行总次数：0

代码可运行

在本篇博客中，我们将深入讲解 Transformer 模型在中英翻译任务中的应用。我们将从最基础的概念开始讲解，一步步带领你理解 Transformer 的各个知识点，并通过实际代码展示如何实现一个简单的中英翻译系统。

一、什么是 Transformer？

Transformer 是一种广泛应用于自然语言处理（NLP）任务的深度学习模型，它由 Vaswani 等人于 2017 年提出。Transformer 模型不依赖传统的循环神经网络（RNN）或长短时记忆网络（LSTM），而是通过“注意力机制”对输入数据进行建模。

注意力机制

注意力机制的核心思想是让模型“关注”输入的某些部分，而非平等对待所有部分。Transformer 中的自注意力机制（Self-Attention）使得每个输入元素都可以参考其他元素的信息，从而能够捕捉到更丰富的上下文信息。

二、Transformer 模型结构

Transformer 模型由编码器（Encoder）和解码器（Decoder）两部分组成，分别用于处理输入序列和生成输出序列。我们首先了解模型的关键组成部分。

1. 自注意力机制（Self-Attention）

自注意力机制帮助模型通过计算输入序列的权重来决定每个词在当前上下文中的重要性。每个词都会生成三个向量：查询（Query）、键（Key）和值（Value）。通过这些向量计算相似度，模型决定哪些词对当前词更重要。

2. 多头注意力（Multi-Head Attention）

为了增强模型的学习能力，Transformer 使用了多头注意力机制。通过并行计算多个“注意力头”，模型可以从不同的角度去“理解”输入数据。

3. 前馈神经网络（Feed-Forward Neural Network）

每一层的输出会通过一个全连接的前馈神经网络进行非线性变换。

4. 位置编码（Positional Encoding）

由于 Transformer 没有递归结构，它无法像 RNN 一样自动捕捉序列的位置信息，因此我们需要使用位置编码来显式地告诉模型每个词在序列中的位置。

5. 残差连接与层归一化（Residual Connection & Layer Normalization）

为了防止深层网络训练困难，Transformer 在每一层的输入和输出之间加入了残差连接。每一层的输出还会进行层归一化，以保证训练的稳定性。

三、使用 PyTorch 实现 Transformer

现在，我们使用 PyTorch 来实现一个简化版的 Transformer 模型，完成中英翻译的任务。首先，我们需要准备数据集。

3.1 数据预处理

我们使用 PyTorch 和 torchtext 库来加载和处理机器翻译数据集，这里我们选择 WMT 数据集（中英翻译）。在这个例子中，我们使用一个小规模的子集进行演示。

pip install torchtext

import torch
import torchtext
from torchtext.datasets import IWSLT
from torchtext.data import Field, BucketIterator

# 1. 定义字段
SRC = Field(tokenize='spacy', tokenizer_language='zh', init_token='<sos>', eos_token='<eos>', lower=True)
TRG = Field(tokenize='spacy', tokenizer_language='en', init_token='<sos>', eos_token='<eos>', lower=True)

# 2. 加载数据
train_data, valid_data, test_data = IWSLT.splits(exts=('.zh', '.en'), fields=(SRC, TRG))

# 3. 构建词汇表
SRC.build_vocab(train_data, min_freq=2)
TRG.build_vocab(train_data, min_freq=2)

# 4. 创建数据迭代器
BATCH_SIZE = 32
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data), 
    batch_size=BATCH_SIZE, 
    device=device)

3.2 Transformer 模型实现

接下来，我们构建 Transformer 模型。模型结构包括：嵌入层、编码器、解码器和最终的线性层。

import torch.nn as nn
import torch.optim as optim

class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, hid_dim, n_layers, n_heads, pf_dim, dropout):
        super().__init__()

        self.encoder = nn.Embedding(input_dim, hid_dim)
        self.decoder = nn.Embedding(output_dim, hid_dim)
        
        self.attn = nn.MultiheadAttention(hid_dim, n_heads, dropout=dropout)
        
        self.fc = nn.Linear(hid_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

        # 位置编码
        self.positional_encoding = nn.Embedding(5000, hid_dim)  # 位置编码的最大长度为5000

    def forward(self, src, trg):
        src = self.encoder(src) + self.positional_encoding(torch.arange(src.size(0), device=src.device))
        trg = self.decoder(trg) + self.positional_encoding(torch.arange(trg.size(0), device=trg.device))

        # 编码器部分
        src = src.transpose(0, 1)
        output, _ = self.attn(src, src, src)

        # 解码器部分
        trg = trg.transpose(0, 1)
        output, _ = self.attn(trg, output, output)

        output = output.transpose(0, 1)
        output = self.fc(output)
        return output

3.3 训练模型

训练过程中，我们使用交叉熵损失函数，并选择 Adam 优化器。

# 参数设置
input_dim = len(SRC.vocab)
output_dim = len(TRG.vocab)
hid_dim = 256
n_layers = 3
n_heads = 8
pf_dim = 512
dropout = 0.1

# 创建模型
model = Transformer(input_dim, output_dim, hid_dim, n_layers, n_heads, pf_dim, dropout).to(device)

# 优化器与损失函数
optimizer = optim.Adam(model.parameters(), lr=0.0005)
criterion = nn.CrossEntropyLoss(ignore_index=TRG.vocab.stoi['<pad>'])

# 训练函数
def train(model, iterator, optimizer, criterion):
    model.train()
    epoch_loss = 0

    for batch in iterator:
        src, trg = batch.src, batch.trg

        optimizer.zero_grad()

        output = model(src, trg[:-1,:])

        # 计算损失
        output_dim = output.shape[-1]
        output = output.view(-1, output_dim)
        trg = trg[1:,:].view(-1)
        
        loss = criterion(output, trg)
        loss.backward()

        optimizer.step()
        epoch_loss += loss.item()

    return epoch_loss / len(iterator)

# 训练模型
for epoch in range(10):
    train_loss = train(model, train_iterator, optimizer, criterion)
    print(f'Epoch {epoch+1}, Train Loss: {train_loss:.3f}')

3.4 使用 GPU 加速训练

在 PyTorch 中，训练模型时可以很容易地使用 GPU 进行加速。只需要确保模型和数据都被转移到 GPU。

# 转移到 GPU（如果有的话）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# 批量数据转移到 GPU
src = src.to(device)
trg = trg.to(device)

3.5 使用 NPU 推理

如果你有 NPU 设备（例如华为昇腾 AI 处理器），你可以使用 MindSpore 或 TensorFlow Lite 来部署 NPU 推理。以下是使用 PyTorch 转化模型到 TensorFlow Lite 的方法：

pip install tf-nightly

将 PyTorch 模型转化为 TensorFlow 格式并进行推理：

import torch
import tensorflow as tf

# 转换 PyTorch 模型为 TensorFlow
import tf2onnx
onnx_model = tf2onnx.convert.from_pytorch(model)

# 将模型保存为 TFLite 格式
converter = tf.lite.TFLiteConverter.from_onnx_model(onnx_model)
tflite_model = converter.convert()

# 保存模型
with open("transformer_model.tflite", "wb") as f:
    f.write(tflite_model)