当前位置：首页 > news >正文

pytorch实现transformer模型

news 2025/7/21 6:02:36

Transformer是一种强大的神经网络架构，可用于处理序列数据，例如自然语言处理任务。在PyTorch中，可以使用torch.nn.Transformer类轻松实现Transformer模型。
以下是一个简单的Transformer模型实现的示例代码，它将一个输入序列转换为一个输出序列，可以用于序列到序列的翻译任务：
示例代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
import mathclass PositionalEncoding(nn.Module):def __init__(self, d_model, dropout=0.1, max_len=5000):super().__init__()self.dropout = nn.Dropout(p=dropout)pe = torch.zeros(max_len, d_model)position = torch.arange(0, max_len, 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)pe = pe.unsqueeze(0).transpose(0, 1)self.register_buffer('pe', pe)def forward(self, x):x = x + self.pe[:x.size(0), :]return self.dropout(x)class TransformerModel(nn.Module):def __init__(self, input_vocab_size, output_vocab_size, d_model, nhead, num_layers, dim_feedforward, dropout=0.1):super(TransformerModel, self).__init__()self.d_model = d_modelself.nhead = nheadself.num_layers = num_layersself.dim_feedforward = dim_feedforwardself.embedding = nn.Embedding(input_vocab_size, d_model)self.pos_encoder = PositionalEncoding(d_model, dropout)encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)self.decoder = nn.Linear(d_model, output_vocab_size)self.init_weights()def init_weights(self):initrange = 0.1self.embedding.weight.data.uniform_(-initrange, initrange)self.decoder.bias.data.zero_()self.decoder.weight.data.uniform_(-initrange, initrange)def forward(self, src, src_mask=None):src = self.embedding(src) * math.sqrt(self.d_model)src = self.pos_encoder(src)output = self.transformer_encoder(src, src_mask)output = self.decoder(output)return output

在上面的代码中，我们定义了一个名为TransformerModel的模型类，它继承自nn.Module。该模型包括以下组件：

nn.Embedding：将输入序列中的每个标记转换为其向量表示。
PositionalEncoding：将序列中每个标记的位置编码为向量。
nn.TransformerEncoder：将编码后的输入序列转换为输出序列。
nn.Linear：将Transformer的输出转换为最终输出序列。
可以根据自己的需求修改TransformerModel类中的超参数，例如输入和输出词汇表大小、嵌入维度、Transformer层数、隐藏层维度等等。使用该模型进行训练时，您需要定义损失函数和优化器，并使用PyTorch的标准训练循环进行训练。

在 Transformer 中，Positional Encoding 的作用是将输入序列中的位置信息嵌入到向量空间中，从而使得每个位置对应的向量是唯一的。这个实现中，Positional Encoding 采用了公式：

$PE(pos,2i)=sin⁡(pos/100002i/dmodel)\text{PE}{(pos, 2i)} = \sin(pos / 10000^{2i/d{\text{model}}})$