Transformer 代码示例

以下是一个简单的 Transformer 模型的 PyTorch 代码示例：

import torch
import torch.nn as nn
import torch.optim as optim

class TransformerModel(nn.Module):
    def __init__(self, input_dim, output_dim, num_layers, hidden_dim, num_heads, dropout):
        super(TransformerModel, self).__init__()

        # 创建嵌入层
        self.embedding = nn.Embedding(input_dim, hidden_dim)

        # 创建 Transformer 编码器
        self.encoder_layer = nn.TransformerEncoderLayer(hidden_dim, num_heads, hidden_dim, dropout)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers)

        # 创建输出层
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, src):
        # 嵌入输入序列
        embedded = self.embedding(src)

        # 将嵌入序列转换为 Transformer 编码器所需的格式
        embedded = embedded.permute(1, 0, 2)

        # 使用 Transformer 编码器编码输入序列
        encoded = self.encoder(embedded)

        # 取出最后一个时间步的输出
        final_output = encoded[-1]

        # 使用全连接层得到最终输出
        output = self.fc(final_output)

        return output

这段代码定义了一个名为 `TransformerModel` 的 PyTorch 模型类，其中包含了一个嵌入层、一个 Transformer 编码器和一个全连接层。在 `__init__` 方法中，我们定义了模型的输入维度、输出维度、编码器的层数、隐藏层维度、注意力头数和 dropout 率。在 `forward` 方法中，我们将输入序列嵌入到嵌入层中，然后将嵌入序列转换为 Transformer 编码器所需的格式，最后使用编码器编码输入序列并取出最后一个时间步的输出。最终输出通过全连接层得到。