pytorch实现transformer模型

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

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class 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_model
        self.nhead = nhead
        self.num_layers = num_layers
        self.dim_feedforward = dim_feedforward

        self.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.1
        self.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 ( p o s , 2 i ) = sin ⁡ ( p o s / 1000 0 2 i / d model ) \text{PE}{(pos, 2i)} = \sin(pos / 10000^{2i/d{\text{model}}}) PE(pos,2i)=sin(pos/100002i/dmodel)

PE ( p o s , 2 i + 1 ) = cos ⁡ ( p o s / 1000 0 2 i / d model ) \text{PE}{(pos, 2i+1)} = \cos(pos / 10000^{2i/d{\text{model}}}) PE(pos,2i+1)=cos(pos/100002i/dmodel)

其中 pos 表示输入序列中的位置,i 表示向量的维度。最终得到的 Positional Encoding 矩阵被添加到输入序列的嵌入向量中。

猜你喜欢

转载自blog.csdn.net/qq_23345187/article/details/129357428
今日推荐