PyTorch 基础学习（10）- Transformer

系列文章：
PyTorch 基础学习（1） - 快速入门
PyTorch 基础学习（2）- 张量 Tensors
PyTorch 基础学习（3） - 张量的数学操作
PyTorch 基础学习（4）- 张量的类型
PyTorch 基础学习（5）- 神经网络
PyTorch 基础学习（6）- 函数API
PyTorch 基础学习（7）- 自动微分
PyTorch 基础学习（8）- 多进程并发
PyTorch 基础学习（9）- 训练优化器

介绍

Transformer模型是近年来在自然语言处理（NLP）领域中非常流行的一种模型架构，尤其是在机器翻译任务中表现出了优异的性能。与传统的循环神经网络（RNN）不同，Transformer模型完全基于注意力机制，避免了序列处理中的长距离依赖问题。本教程将通过一个简单的实例，详细讲解如何在PyTorch中实现一个基于Transformer的机器翻译模型。

Transformer的原理简介

Transformer模型由Vaswani等人在2017年提出，其核心思想是利用注意力机制来捕捉输入序列中的长程依赖关系。模型主要包括两个模块：编码器（Encoder）和解码器（Decoder）。每个模块由多个层（Layer）堆叠而成，每一层又包含多个子层（Sub-layer），如自注意力机制（Self-Attention）、前馈神经网络（Feed-Forward Neural Network）等。

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

自注意力机制是Transformer的核心，主要用于计算输入序列中各元素之间的相互依赖关系。通过自注意力机制，模型可以在每一步中考虑到整个序列的信息，而不是仅仅依赖于固定的上下文窗口。

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

多头注意力机制是对自注意力机制的扩展，通过引入多个注意力头（Attention Heads），模型可以在不同的子空间中独立地计算注意力，从而捕捉到输入序列中更多的特征。

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

在每个编码器和解码器层中，注意力机制后接一个前馈神经网络。该网络在每个时间步上独立应用于序列中的每一个位置。

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

为了缓解梯度消失的问题，Transformer模型在每个子层之间使用了残差连接，并在每个子层后使用层归一化。

实例代码及讲解

下面我们将通过一个简单的示例代码，详细讲解如何在PyTorch中实现一个基于Transformer的句子推理。

1. 导入必要的库

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence
import numpy as np

2. 定义数据集类

class TranslationDataset(Dataset):
    def __init__(self, source_sentences, target_sentences, src_vocab, tgt_vocab):
        self.source_sentences = source_sentences
        self.target_sentences = target_sentences
        self.src_vocab = src_vocab
        self.tgt_vocab = tgt_vocab

    def __len__(self):
        return len(self.source_sentences)

    def __getitem__(self, idx):
        src = [self.src_vocab[word] for word in self.source_sentences[idx].split()]
        tgt = [self.tgt_vocab[word] for word in self.target_sentences[idx].split()]
        return torch.tensor(src), torch.tensor(tgt)

• TranslationDataset类继承自Dataset，用于处理机器翻译任务中的数据集。
• __getitem__方法根据索引idx返回源句子和目标句子的张量表示。

3. 定义collate_fn函数

def collate_fn(batch):
    src_batch, tgt_batch = zip(*batch)
    src_batch = pad_sequence(src_batch, padding_value=0, batch_first=True)
    tgt_batch = pad_sequence(tgt_batch, padding_value=0, batch_first=True)
    return src_batch, tgt_batch

• collate_fn函数用于将一个批次的数据进行填充，使得每个批次的源句子和目标句子长度一致。

4. 定义Transformer模型

class TransformerModel(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6,
                 dim_feedforward=2048, dropout=0.1):
        super(TransformerModel, self).__init__()
        self.embedding_src = nn.Embedding(src_vocab_size, d_model)
        self.embedding_tgt = nn.Embedding(tgt_vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward,
                                          dropout)
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
        self.d_model = d_model

    def forward(self, src, tgt):
        src = self.embedding_src(src) * np.sqrt(self.d_model)
        tgt = self.embedding_tgt(tgt) * np.sqrt(self.d_model)
        src = src.permute(1, 0, 2)
        tgt = tgt.permute(1, 0, 2)
        output = self.transformer(src, tgt)
        output = self.fc_out(output)
        return output

    def generate(self, src, max_len, sos_idx):
        self.eval()
        src = self.embedding_src(src) * np.sqrt(self.d_model)
        src = src.permute(1, 0, 2)  # [sequence_length, batch_size, d_model]
        memory = self.transformer.encoder(src)

        # 初始化解码器输入，开始标记
        ys = torch.ones(1, 1).fill_(sos_idx).type(torch.long).to(src.device)

        for i in range(max_len - 1):
            tgt = self.embedding_tgt(ys) * np.sqrt(self.d_model)
            tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(0)).to(src.device)

            out = self.transformer.decoder(tgt, memory, tgt_mask=tgt_mask)

            out = self.fc_out(out)
            prob = out[-1, :, :].max(dim=-1)[1]
            ys = torch.cat([ys, prob.unsqueeze(0)], dim=0)
            if prob == 2:  # <eos> token index
                break
        return ys.transpose(0, 1)

• TransformerModel类继承自nn.Module，封装了Transformer模型。
• forward方法定义了模型的前向传播逻辑，包括对源句子和目标句子进行嵌入、通过Transformer层处理，以及通过线性层输出预测结果。
• generate方法用于推理，生成翻译结果。

5. 训练和评估函数

def train(model, dataloader, optimizer, criterion, num_epochs=10):
    model.train()
    for epoch in range(num_epochs):
        epoch_loss = 0
        for src, tgt in dataloader:
            tgt_input = tgt[:, :-1]
            tgt_output = tgt[:, 1:]

            optimizer.zero_grad()

            output = model(src, tgt_input)
            output = output.view(-1, output.shape[-1])
            tgt_output = tgt_output.reshape(-1)

            loss = criterion(output, tgt_output)
            loss.backward()

            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()

            epoch_loss += loss.item()
        print(f'Epoch {epoch + 1}, Loss: {epoch_loss / len(dataloader)}')

def evaluate(model, dataloader, criterion):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for src, tgt in dataloader:
            tgt_input = tgt[:, :-1]
            tgt_output = tgt[:, 1:]

            output = model(src, tgt_input)
            output = output.view(-1, output.shape[-1])
            tgt_output = tgt_output.reshape(-1)

            loss = criterion(output, tgt_output)
            total_loss += loss.item()

    print(f'Evaluation Loss: {total_loss / len(dataloader)}')

• train函数用于训练模型，逐批处理数据，计算损失，并更新模型参数。
• evaluate函数用于评估模型的性能，计算整个数据集的平均损失。

6. 推理函数

def inference(model, src_sentence, src_vocab, tgt_vocab, max_len=2):
    model.eval()
    src_indexes = [src_vocab[word] for word in src_sentence.split()]
    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(next(model.parameters()).device)  # 确保是 LongTensor 类型
    sos_idx = tgt_vocab["<sos>"]
    generated_tensor = model.generate(src_tensor, max_len, sos_idx)
    generated_sentence = ' '.join([list(tgt_vocab.keys())[i] for i in generated_tensor.squeeze().tolist()])
    return generated_sentence

• inference函数用于对单个句子进行翻译，生成对应的目标句子。

7. 运行示例

if __name__ == "__main__":
    # 假设我们有一个简单的词汇表和句子对
    vocab = {
        "<pad>": 0,
        "<sos>": 1,
        "<eos>": 2,
        "hello": 3,
        "world": 4,
        "good": 5,
        "morning": 6,
        "night": 7,
        "how": 8,
        "are": 9,
        "you": 10,
        "today": 11,
        "friend": 12,
        "goodbye": 13,
        "see": 14,
        "take": 15,
        "care": 16,
        "welcome": 17,
        "back": 18
    }

    sentences = [
        "hello world",
        "good morning",
        "goodbye friend",
        "see you",
        "take care",
        "welcome back",
    ]

    src_vocab = vocab
    tgt_vocab = vocab
    source_sentences = sentences
    target_sentences = sentences

    # 创建数据集和数据加载器
    dataset = TranslationDataset(source_sentences, target_sentences, src_vocab, tgt_vocab)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)

    # 模型初始化
    model = TransformerModel(len(src_vocab), len(tgt_vocab))
    optimizer = optim.Adam(model.parameters(), lr=0.0001)
    criterion = nn.CrossEntropyLoss(ignore_index=0)

    # 训练模型
    train(model, dataloader, optimizer, criterion, num_epochs=20)

    # 评估模型
    evaluate(model, dataloader, criterion)

    # 推理测试
    test_sentence = "hello"
    translated_sentence = inference(model, test_sentence, src_vocab, tgt_vocab)
    print(f"Input: {test_sentence}")
    print(f"Output: {translated_sentence}")

    test_sentence = "see"
    translated_sentence = inference(model, test_sentence, src_vocab, tgt_vocab)
    print(f"Input: {test_sentence}")
    print(f"Output: {translated_sentence}")

    test_sentence = "welcome"
    translated_sentence = inference(model, test_sentence, src_vocab, tgt_vocab)
    print(f"Input: {test_sentence}")
    print(f"Output: {translated_sentence}")

在这个运行示例中，我们首先定义了一个简单的词汇表和句子对，然后创建数据集和数据加载器。接下来，我们初始化Transformer模型，设置优化器和损失函数，训练模型并进行评估。最后，通过推理函数对一些输入句子进行翻译，并输出结果。

完整代码实例

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torch.nn.utils.rnn import pad_sequence
import numpy as np

# 定义数据集类，用于加载源语言和目标语言的句子
class TranslationDataset(Dataset):
    def __init__(self, source_sentences, target_sentences, src_vocab, tgt_vocab):
        self.source_sentences = source_sentences  # 源语言句子列表
        self.target_sentences = target_sentences  # 目标语言句子列表
        self.src_vocab = src_vocab  # 源语言词汇表
        self.tgt_vocab = tgt_vocab  # 目标语言词汇表

    def __len__(self):
        return len(self.source_sentences)  # 返回数据集中句子的数量

    def __getitem__(self, idx):
        # 将源语言和目标语言的句子转换为词汇表中的索引
        src = [self.src_vocab[word] for word in self.source_sentences[idx].split()]
        tgt = [self.tgt_vocab[word] for word in self.target_sentences[idx].split()]
        return torch.tensor(src), torch.tensor(tgt)  # 返回源句子和目标句子的索引张量

# 定义collate_fn函数，用于在批处理中对序列进行填充
def collate_fn(batch):
    src_batch, tgt_batch = zip(*batch)  # 将批次中的源和目标句子分开
    src_batch = pad_sequence(src_batch, padding_value=0, batch_first=True)  # 对源句子进行填充
    tgt_batch = pad_sequence(tgt_batch, padding_value=0, batch_first=True)  # 对目标句子进行填充
    return src_batch, tgt_batch  # 返回填充后的源和目标句子张量

# 定义Transformer模型
class TransformerModel(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6,
                 dim_feedforward=2048, dropout=0.1):
        super(TransformerModel, self).__init__()
        # 定义源语言和目标语言的嵌入层
        self.embedding_src = nn.Embedding(src_vocab_size, d_model)
        self.embedding_tgt = nn.Embedding(tgt_vocab_size, d_model)
        # 定义Transformer模型
        self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward,
                                          dropout)
        # 定义输出的全连接层，将Transformer的输出转换为词汇表中的分布
        self.fc_out = nn.Linear(d_model, tgt_vocab_size)
        self.d_model = d_model  # d_model是嵌入向量的维度

    def forward(self, src, tgt):
        # 将源语言和目标语言的索引转换为嵌入向量，并进行缩放
        src = self.embedding_src(src) * np.sqrt(self.d_model)
        tgt = self.embedding_tgt(tgt) * np.sqrt(self.d_model)
        # 调整维度以适应Transformer输入的要求
        src = src.permute(1, 0, 2)
        tgt = tgt.permute(1, 0, 2)
        # 将源语言和目标语言嵌入输入到Transformer中
        output = self.transformer(src, tgt)
        # 使用全连接层将Transformer的输出转换为目标词汇表中的分布
        output = self.fc_out(output)
        return output

    def generate(self, src, max_len, sos_idx):
        self.eval()  # 设置模型为评估模式
        # 对源语言进行嵌入并缩放
        src = self.embedding_src(src) * np.sqrt(self.d_model)
        src = src.permute(1, 0, 2)  # 调整维度
        memory = self.transformer.encoder(src)  # 通过编码器获取源语言的记忆表示

        # 初始化解码器输入，使用<start of sequence>标记
        ys = torch.ones(1, 1).fill_(sos_idx).type(torch.long).to(src.device)

        for i in range(max_len - 1):
            # 对目标语言进行嵌入并缩放
            tgt = self.embedding_tgt(ys) * np.sqrt(self.d_model)

            # 生成用于掩码的下三角矩阵，以确保模型不能看到未来的词
            tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(0)).to(src.device)

            # 使用Transformer解码器生成输出
            out = self.transformer.decoder(tgt, memory, tgt_mask=tgt_mask)

            out = self.fc_out(out)  # 通过全连接层生成词汇表的分布
            prob = out[-1, :, :].max(dim=-1)[1]  # 选择概率最大的词作为输出
            # 将生成的词拼接到解码器的输入中
            ys = torch.cat([ys, prob.unsqueeze(0)], dim=0)
            if prob == 2:  # 如果生成了<end of sequence>标记，则停止生成
                break
        return ys.transpose(0, 1)  # 返回生成的序列

# 训练函数
def train(model, dataloader, optimizer, criterion, num_epochs=10):
    model.train()  # 设置模型为训练模式
    for epoch in range(num_epochs):
        epoch_loss = 0  # 记录每个epoch的损失
        for src, tgt in dataloader:
            tgt_input = tgt[:, :-1]  # 获取目标句子中除了最后一个词的部分作为输入
            tgt_output = tgt[:, 1:]  # 获取目标句子中除了第一个词的部分作为输出

            optimizer.zero_grad()  # 清零梯度

            output = model(src, tgt_input)  # 前向传播计算输出
            output = output.view(-1, output.shape[-1])  # 将输出展平为2D张量
            tgt_output = tgt_output.reshape(-1)  # 将目标输出展平为1D张量

            loss = criterion(output, tgt_output)  # 计算损失
            loss.backward()  # 反向传播计算梯度

            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 对梯度进行裁剪以防止梯度爆炸

            optimizer.step()  # 更新模型参数

            epoch_loss += loss.item()  # 累加损失
        print(f'Epoch {epoch + 1}, Loss: {epoch_loss / len(dataloader)}')  # 输出每个epoch的平均损失

# 评估函数
def evaluate(model, dataloader, criterion):
    model.eval()  # 设置模型为评估模式
    total_loss = 0  # 记录总损失
    with torch.no_grad():  # 在评估时不需要计算梯度
        for src, tgt in dataloader:
            tgt_input = tgt[:, :-1]  # 获取目标句子中除了最后一个词的部分作为输入
            tgt_output = tgt[:, 1:]  # 获取目标句子中除了第一个词的部分作为输出

            output = model(src, tgt_input)  # 前向传播计算输出
            output = output.view(-1, output.shape[-1])  # 将输出展平为2D张量
            tgt_output = tgt_output.reshape(-1)  # 将目标输出展平为1D张量

            loss = criterion(output, tgt_output)  # 计算损失
            total_loss += loss.item()  # 累加损失

    print(f'Evaluation Loss: {total_loss / len(dataloader)}')  # 输出平均评估损失

# 推理函数，用于在模型训练完毕后进行句子翻译
def inference(model, src_sentence, src_vocab, tgt_vocab, max_len=2):
    model.eval()  # 设置模型为评估模式
    # 将源语言句子转换为索引序列
    src_indexes = [src_vocab[word] for word in src_sentence.split()]
    # 将索引序列转换为张量，并添加批次维度
    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(next(model.parameters()).device)
    sos_idx = tgt_vocab["<sos>"]  # 获取<sos>标记的索引
    # 使用模型生成目标语言的句子
    generated_tensor = model.generate(src_tensor, max_len, sos_idx)
    # 将生成的索引序列转换为词语序列
    generated_sentence = ' '.join([list(tgt_vocab.keys())[i] for i in generated_tensor.squeeze().tolist()])
    return generated_sentence  # 返回生成的句子

# 示例运行
if __name__ == "__main__":
    # 假设我们有一个简单的词汇表和句子对
    vocab = {
        "<pad>": 0,
        "<sos>": 1,
        "<eos>": 2,
        "hello": 3,
        "world": 4,
        "good": 5,
        "morning": 6,
        "night": 7,
        "how": 8,
        "are": 9,
        "you": 10,
        "today": 11,
        "friend": 12,
        "goodbye": 13,
        "see": 14,
        "take": 15,
        "care": 16,
        "welcome": 17,
        "back": 18
    }

    sentences = [
        "hello world",
        "good morning",
        "goodbye friend",
        "see you",
        "take care",
        "welcome back",
    ]

    src_vocab = vocab  # 源语言词汇表
    tgt_vocab = vocab  # 目标语言词汇表
    source_sentences = sentences  # 源语言句子列表
    target_sentences = sentences  # 目标语言句子列表

    # 创建数据集和数据加载器
    dataset = TranslationDataset(source_sentences, target_sentences, src_vocab, tgt_vocab)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)

    # 模型初始化
    model = TransformerModel(len(src_vocab), len(tgt_vocab))
    optimizer = optim.Adam(model.parameters(), lr=0.0001)
    criterion = nn.CrossEntropyLoss(ignore_index=0)  # 使用交叉熵损失函数，忽略填充标记的损失

    # 训练模型
    train(model, dataloader, optimizer, criterion, num_epochs=20)

    # 评估模型
    evaluate(model, dataloader, criterion)

    # 推理测试
    test_sentence = "hello"
    translated_sentence = inference(model, test_sentence, src_vocab, tgt_vocab)
    print(f"Input: {test_sentence}")
    print(f"Output: {translated_sentence}")

    test_sentence = "see"
    translated_sentence = inference(model, test_sentence, src_vocab, tgt_vocab)
    print(f"Input: {test_sentence}")
    print(f"Output: {translated_sentence}")

    test_sentence = "welcome"
    translated_sentence = inference(model, test_sentence, src_vocab, tgt_vocab)
    print(f"Input: {test_sentence}")
    print(f"Output: {translated_sentence}")

运行结果：

......
Epoch 18, Loss: 0.0005644524741607407
Epoch 19, Loss: 0.0005254073378940424
Epoch 20, Loss: 0.0004640306190898021
Evaluation Loss: 0.00014784792438149452
Input: hello
Output: <sos> world
Input: see
Output: <sos> you
Input: welcome
Output: <sos> back

总结

通过这个教程，我们从理论到实践，详细讲解了Transformer模型的基本原理，并展示了如何使用PyTorch实现一个简单的机器推理模型。虽然这个示例中的模型和数据集都非常简化，但它为进一步学习和研究更复杂的NLP任务打下了基础。希望通过这个教程，你能够对Transformer模型有更深入的理解，并能够在自己的项目中灵活应用。