PyTorch实战Transformer模型：从零开始构建

最新推荐文章于 2024-07-24 00:12:22 发布

DOUZHAOPENG1113

最新推荐文章于 2024-07-24 00:12:22 发布

阅读量137

点赞数

文章标签： pytorch transformer 人工智能

原文链接：https://acg.stackbox.cn

版权

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 采用了公式：

(

)

sin

⁡

(