Pytorch 实现 Transformer 模型

Python使用第三方库Pytorch 实现 Transformer 模型

PyTorch的Transformer模型是基于神经网络模块实现的，它采用了编码器-解码器的结构，其中编码器和解码器都由多层的自注意力机制和前馈神经网络组成，使用PyTorch实现Transformer模型请确保已经安装对应的torch库，如下是一个实现Transformer的代码示例：

import torch

class Transformer(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_heads, num_layers):
        # 继承 torch.nn.Module 中的属性
        super(Transformer, self).__init__()
        # 嵌入层
        self.embedding = torch.nn.Module.Embedding(input_dim, hidden_dim)
        # 位置编码层
        self.positional_encoding = PositionalEncoding(hidden_dim)
        # 编码器
        self.encoder = torch.nn.Module.TransformerEncoder(
            torch.nn.Module.TransformerEncoderLayer(hidden_dim, num_heads), num_layers)
        # 全连接层
        self.fc = torch.nn.Module.Linear(hidden_dim, output_dim)

    # 实现模型的向前计算并返回预测结果
    def forward(self, x):
        embedded = self.embedding(x)
        embedded = self.positional_encoding(embedded)
        encoded = self.encoder(embedded)
        output = self.fc(encoded)
        # 返回模型的预测结果
        return output

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        # 初始化位置编码矩阵 pe
        pe = torch.zeros(max_len, d_model)
        # 生成张量
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # 使用 exp 函数调整编码器
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(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):
        return x + self.pe[:x.size(0), :]

# 输入维度 隐层维度 输出维度 注意力头数 Transformer 层数
input_dim, hidden_dim, output_dim, num_heads, num_layers = 100, 128, 10, 4, 2
model = Transformer(input_dim, hidden_dim, output_dim, num_heads, num_layers)
batch_size, sequence_length = 32, 10
# 创建随机输入序列
input_data = torch.randint(0, input_dim, (batch_size, sequence_length))
print(model(input_data))

上述代码实现了Transformer模型，首先定义了一个Transformer类继承自torch.nn.Module类，在构造函数__init__中，进行了模型的初始化操作，其中self.embedding是一个嵌入层，用于将输入的离散化整数序列转换为连续的向量表示，self.positional_encoding是一个位置编码层，用于编码输入序列中每个元素的位置信息，self.encoder是一个Transformer编码器，用于对输入序列进行编码和特征提取，self.fc是一个全连接层，用于将Transformer编码器的输出映射到最终的输出维度，在forward方法中，实现了整个Transformer模型的前向传播计算，输入数据x首先经过嵌入层和位置编码层，然后输入Transformer编码器进行编码。最后，通过全连接层将编码结果映射到最终的输出维度，并返回输出。
接下来，定义了一个PositionalEncoding类，也是继承自nn.Module类。在构造函数__init__中，生成了一个位置编码矩阵pe，其大小为 max_len和d_model，其中max_len表示序列的最大长度，d_model表示隐藏维度。通过一定的数学计算，为每个位置和每个维度生成不同的正弦sin和余弦cos编码值，并进行了适当的归一化处理。最后，将位置编码矩阵pe注册为模型的缓冲区。