Transformer PyTorch 教程

Transformer 架构自 2017 年提出以来，已经成为自然语言处理领域的明星模型。本文将介绍如何使用 PyTorch 实现一个基本的 Transformer 模型。

1. 简介

Transformer 是一种基于自注意力机制的深度神经网络模型，常用于序列到序列的学习任务，如机器翻译、文本摘要等。

2. PyTorch Transformer 库

PyTorch 提供了一个名为 torch.nn.Transformer 的库，可以方便地实现 Transformer 模型。

2.1. 模型结构

Transformer 模型主要由以下几个部分组成：

编码器：将输入序列编码成向量表示。
解码器：将编码器的输出解码成输出序列。
注意力机制：通过注意力机制来学习输入序列中不同位置之间的关系。

2.2. 示例代码

以下是一个使用 PyTorch 实现的简单 Transformer 模型示例：

import torch
import torch.nn as nn

class TransformerModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.transformer = nn.Transformer(d_model, nhead)
        self.fc = nn.Linear(d_model, vocab_size)

    def forward(self, src, tgt):
        src = self.embedding(src)
        tgt = self.embedding(tgt)
        output = self.transformer(src, tgt)
        output = self.fc(output)
        return output

3. 损失函数与优化器

在训练 Transformer 模型时，常用的损失函数是交叉熵损失（CrossEntropyLoss），优化器可以使用 Adam 或 SGD。

import torch.optim as optim

model = TransformerModel(vocab_size=10000, d_model=512, nhead=8)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 假设有一个输入序列和目标序列
input_seq = torch.tensor([[1, 2, 3], [4, 5, 6]])
target_seq = torch.tensor([[2, 3, 4], [5, 6, 7]])

optimizer.zero_grad()
output = model(input_seq, target_seq)
loss = criterion(output.view(-1, vocab_size), target_seq.view(-1))
loss.backward()
optimizer.step()

4. 扩展阅读

想要深入了解 Transformer 和 PyTorch 的朋友们可以参考以下链接：