用 Pytorch 训练一个 Transformer模型

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

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

  1. 准备环境。
  2. 然后就是跟据架构来定义每一层,包括Embedding、Position Encoding、多头注意力、 网络层。
  3. 准备Encoder。
  4. 准备Decoder。
  5. 运行 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)]

多头注意力

  1. 初始化model 维度,头数,每个头的维度。
  2. 计算是在 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
<think>我们正在处理用户关于使用PyTorch编写Transformer模型的问题。用户之前询问了关于单线程模拟信号量的问题,但当前问题转向了Transformer模型。根据引用内容,用户希望获取使用PyTorch编写Transformer模型的教程。 引用[1]和[2]强调了PyTorch在构建Transformer模型方面的优势,引用[3]提供了一个构建Transformer模型的代码片段。因此,我们将基于这些引用,并结合PyTorch官方文档和常见实践,提供一个详细的构建Transformer模型的教程。 构建Transformer模型的关键组件包括: 1. 嵌入层(Embeddings) 2. 位置编码(PositionalEncoding) 3. 多头注意力(MultiHeadedAttention) 4. 前馈网络(PositionwiseFeedForward) 5. 编码器层(EncoderLayer)和解码器层(DecoderLayer) 6. 编码器(Encoder)和解码器(Decoder) 7. 整个Transformer模型(EncoderDecoder) 我们将按照以下步骤进行: 1. 导入必要的库 2. 实现嵌入层和位置编码 3. 实现多头注意力机制 4. 实现前馈网络 5. 实现编码器层和解码器层 6. 实现编码器和解码器 7. 组合成完整的Transformer模型 8. 初始化模型参数 注意:由于引用[3]中提供了一些代码,我们将参考它,但会进行适当调整和解释,确保代码可运行和易于理解。 下面是一个详细的实现教程: </think>### 使用PyTorch构建Transformer模型的完整指南 以下是基于PyTorch实现Transformer模型的关键步骤和完整代码示例,参考了PyTorch官方实现和经典论文《Attention is All You Need》的设计[^2][^3]: #### 1. 导入基础库 ```python import torch import torch.nn as nn import torch.nn.functional as F import math import copy ``` #### 2. 核心组件实现 ##### 多头注意力机制 ($MultiHeadAttention$) ```python class MultiHeadedAttention(nn.Module): def __init__(self, h, d_model, dropout=0.1): super().__init__() assert d_model % h == 0 self.d_k = d_model // h # 每个头的维度 self.h = h # 头数 self.linears = nn.ModuleList([nn.Linear(d_model, d_model) for _ in range(4)]) self.dropout = nn.Dropout(p=dropout) def forward(self, query, key, value, mask=None): batch_size = query.size(0) # 线性变换并分割多头 query, key, value = [ l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linears, (query, key, value)) ] # 计算缩放点积注意力 scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) p_attn = F.softmax(scores, dim=-1) p_attn = self.dropout(p_attn) x = torch.matmul(p_attn, value) # 合并多头结果 x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) return self.linears[-1](x) ``` ##### 位置前馈网络 ($PositionwiseFeedForward$) ```python class PositionwiseFeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout=0.1): super().__init__() self.w1 = nn.Linear(d_model, d_ff) self.w2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.w2(self.dropout(F.relu(self.w1(x)))) ``` ##### 位置编码 ($PositionalEncoding$) ```python class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout, max_len=5000): super().__init__() self.dropout = nn.Dropout(p=dropout) pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) pe = pe.unsqueeze(0) self.register_buffer('pe', pe) def forward(self, x): x = x + self.pe[:, :x.size(1)] return self.dropout(x) ``` #### 3. 构建编码器-解码器结构 ```python def clones(module, N): "生成N个相同的层" return nn.ModuleList([copy.deepcopy(module) for _ in range(N)]) class Encoder(nn.Module): def __init__(self, layer, N): super().__init__() self.layers = clones(layer, N) self.norm = nn.LayerNorm(layer.size) def forward(self, x, mask): for layer in self.layers: x = layer(x, mask) return self.norm(x) class Decoder(nn.Module): def __init__(self, layer, N): super().__init__() self.layers = clones(layer, N) self.norm = nn.LayerNorm(layer.size) def forward(self, x, memory, src_mask, tgt_mask): for layer in self.layers: x = layer(x, memory, src_mask, tgt_mask) return self.norm(x) ``` #### 4. 完整的Transformer模型 ```python class Transformer(nn.Module): def __init__(self, src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1): super().__init__() c = copy.deepcopy attn = MultiHeadedAttention(h, d_model) ff = PositionwiseFeedForward(d_model, d_ff, dropout) position = PositionalEncoding(d_model, dropout) # 嵌入层 self.src_embed = nn.Sequential( nn.Embedding(src_vocab, d_model), c(position) ) self.tgt_embed = nn.Sequential( nn.Embedding(tgt_vocab, d_model), c(position) ) # 编码器-解码器结构 self.encoder = Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N) self.decoder = Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N) # 输出层 self.generator = nn.Linear(d_model, tgt_vocab) # Xavier初始化 for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def forward(self, src, tgt, src_mask, tgt_mask): "处理源序列和目标序列" memory = self.encode(src, src_mask) return self.decode(memory, src_mask, tgt, tgt_mask) def encode(self, src, src_mask): return self.encoder(self.src_embed(src), src_mask) def decode(self, memory, src_mask, tgt, tgt_mask): return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask) ``` #### 5. 应用示例(机器翻译) ```python # 参数设置 src_vocab_size = 10000 # 源语言词汇表大小 tgt_vocab_size = 8000 # 目标语言词汇表大小 d_model = 512 N = 6 # 编码器/解码器层数 # 实例化模型 model = Transformer(src_vocab_size, tgt_vocab_size, N, d_model) # 输入数据 (批大小=32, 序列长度=100) src = torch.randint(0, src_vocab_size, (32, 100)) tgt = torch.randint(0, tgt_vocab_size, (32, 100)) # 创建掩码 src_mask = torch.ones(32, 1, 100) tgt_mask = torch.tril(torch.ones(100, 100)).unsqueeze(0) # 前向传播 output = model(src, tgt, src_mask, tgt_mask) print(output.shape) # torch.Size([32, 100, 8000]) ``` ### 关键公式说明 1. **缩放点积注意力**: $$ \text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ 2. **位置编码**: $$ PE_{(pos,2i)} = \sin\left(pos / 10000^{2i/d_{\text{model}}}\right) \\ PE_{(pos,2i+1)} = \cos\left(pos / 10000^{2i/d_{\text{model}}}\right) $$ ### 训练优化技巧 1. **学习率调度**: ```python optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step: min(step**(-0.5), step * 4000**(-1.5)) ) ``` 2. **标签平滑**: ```python criterion = nn.CrossEntropyLoss(label_smoothing=0.1) ``` ### 实际应用场景 1. **机器翻译**:实现多语言自动翻译系统[^2] 2. **文本摘要**:生成文章关键摘要 3. **对话系统**:构建智能聊天机器人 4. **代码生成**:根据描述自动生成代码
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值