使用 PyTorch 实现简单 Transformer 模型的示例（含代码）

使用 PyTorch 实现简单 Transformer 模型的示例。这个实现包括了模型的构建以及如何使用它进行前向传播。为了简化，我们将使用 PyTorch 自带的 nn.Transformer 模块来构建 Transformer 模型，而不是从头实现每个组件。

完整实例免费下载链接

1. 安装依赖

确保你已经安装了 PyTorch。如果没有安装，可以使用以下命令进行安装：

pip install torch

2. Transformer 模型实现

下面的代码演示了如何使用 PyTorch 的 nn.Transformer 模块来构建一个简单的 Transformer 模型。

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

# 参数设置
d_model = 512       # 特征维度
nhead = 8           # 注意力头数
num_encoder_layers = 6
num_decoder_layers = 6
dim_feedforward = 2048
dropout = 0.1

# 定义 Transformer 模型
class TransformerModel(nn.Module):
    def __init__(self, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout):
        super(TransformerModel, self).__init__()
        self.transformer = nn.Transformer(
            d_model=d_model,
            nhead=nhead,
            num_encoder_layers=num_encoder_layers,
            num_decoder_layers=num_decoder_layers,
            dim_feedforward=dim_feedforward,
            dropout=dropout
        )
        self.src_embed = nn.Embedding(1000, d_model)  # 词汇表大小为1000
        self.tgt_embed = nn.Embedding(1000, d_model)
        self.fc_out = nn.Linear(d_model, 1000)  # 输出词汇表大小为1000

    def forward(self, src, tgt):
        src = self.src_embed(src)
        tgt = self.tgt_embed(tgt)
        output = self.transformer(src, tgt)
        output = self.fc_out(output)
        return output

# 模型实例化
model = TransformerModel(
    d_model=d_model,
    nhead=nhead,
    num_encoder_layers=num_encoder_layers,
    num_decoder_layers=num_decoder_layers,
    dim_feedforward=dim_feedforward,
    dropout=dropout
)

# 打印模型结构
print(model)

# 模拟数据
src = torch.randint(0, 1000, (10, 32))  # (序列长度, 批量大小)
tgt = torch.randint(0, 1000, (20, 32))  # (序列长度, 批量大小)

# 前向传播
output = model(src, tgt)
print(output.shape)

3. 解释

• 模型构建：TransformerModel 类封装了一个 PyTorch 的 nn.Transformer 模块。我们为输入和目标序列定义了嵌入层 src_embed 和 tgt_embed。最后通过 fc_out 线性层将 Transformer 的输出映射到词汇表大小的维度上。

• 前向传播：在 forward 方法中，我们将输入和目标序列通过嵌入层转换为适合 Transformer 的输入格式，然后传入 nn.Transformer 模块进行处理，最后通过 fc_out 层获取模型的最终输出。

• 模拟数据：我们生成了一些随机的输入数据 src 和 tgt，并通过模型进行前向传播以获得输出。

4. 训练模型

要训练模型，你需要定义损失函数和优化器，并编写训练循环。以下是一个基本的训练循环示例：

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练循环示例
for epoch in range(10):  # 训练10个周期
    model.train()
    optimizer.zero_grad()
    
    # 前向传播
    output = model(src, tgt)
    
    # 计算损失
    loss = criterion(output.view(-1, 1000), tgt.view(-1))  # 视图调整以适应 CrossEntropyLoss
    loss.backward()
    
    # 更新参数
    optimizer.step()
    
    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

这个代码展示了一个简单的训练循环，其中 src 和 tgt 是你的训练数据，1000 是词汇表的大小，你需要根据实际情况调整。