用 Pytorch 训练一个 Transformer模型_pytorch transformer 训练-优快云博客

本文详细介绍了如何在Pytorch中构建Transformer模型，包括环境准备、各层结构（如Embedding、PositionalEncoding、多头注意力和网络层）、Encoder和Decoder的设计，以及模型的训练和评估过程。作者通过示例展示了从安装环境到模型训练的完整流程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

昨天说了一下Transformer架构，今天我们来看看怎么 Pytorch 训练一个Transormer模型，真实训练一个模型是个庞大工程，准备数据、准备硬件等等，我只是做一个简单的实现。因为只是做实验，本地用 CPU 也可以运行。
本文包含以下几部分：

准备环境。
然后就是跟据架构来定义每一层，包括Embedding、Position Encoding、多头注意力、网络层。
准备Encoder。
准备Decoder。
运行 Transformer，包括训练和评估。

安装Pytorch 环境

!pip3 install torch torchvision torchaudio

引入所需工具类库

引入需要的类库，pytorch 是强大的训练框架，深度学习中需要的一些函数和基本功能都已经实现。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import math
import copy

Position Embedding

生成位置信息。

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_seq_length):
        super(PositionalEncoding, self).__init__()
        
        pe = torch.zeros(max_seq_length, d_model)
        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

多头注意力

初始化model 维度，头数，每个头的维度。
计算是在 forward 这个方法，主要看这个方法。

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        # Ensure that the model dimension (d_model) is divisible by the number of heads
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        
        # Initialize dimensions
        self.d_model = d_model # Model's dimension
        self.num_heads = num_heads # Number of attention heads
        self.d_k = d_model // num_heads # Dimensio