transformer 时序预测 pytorch

原创

mob64ca12e676c8 2024-12-22 03:30:43 ©著作权

©著作权归作者所有：来自51CTO博客作者mob64ca12e676c8的原创作品，请联系作者获取转载授权，否则将追究法律责任

Transformer 时序预测的 PyTorch 实现

最近，Transformer 模型因其在自然语言处理（NLP）和其他序列数据任务中的出色表现而受到广泛关注。虽然最初是为文本生成和翻译设计的，但其强大的特性在时序预测中同样适用。本文将介绍如何使用 PyTorch 来实现一个简单的 Transformer 时序预测模型。

Transformer 概述

Transformer 模型利用自注意力机制来捕捉序列数据中的长期依赖关系，而传统的 RNN 和 LSTM 方法可能在处理长序列时面临困难。该模型主要由编码器和解码器两部分组成。对于时序预测任务，我们可以仅使用编码器部分来预测未来的时间步。

序列图示例

为了更好地理解 Transformer 是如何处理序列的，下面是一个简单的序列图示例：

sequenceDiagram
    participant Input as 输入序列
    participant Encoder as Transformer 编码器
    participant Output as 预测输出
    Input->>Encoder: 传递输入序列
    Encoder->>Output: 生成预测结果

安装 PyTorch

要使用 PyTorch，我们首先需要确保安装了 PyTorch。你可以通过以下命令进行安装：

pip install torch torchvision

实现 Transformer 时序预测

在下面的代码示例中，我们将构建一个简单的 Transformer 时序预测模型。

创建一个 TransformerModel 类，包含编码器和前向传播函数：

import torch
import torch.nn as nn

class TransformerModel(nn.Module):
    def __init__(self, n_features, n_heads, n_layers, n_output):
        super(TransformerModel, self).__init__()
        self.transformer = nn.Transformer(
            d_model=n_features,
            nhead=n_heads,
            num_encoder_layers=n_layers,
        )
        self.fc = nn.Linear(n_features, n_output)

    def forward(self, x):
        # x shape: (sequence_length, batch_size, n_features)
        x = self.transformer(x)
        # 只取最后的时间步作为预测输出
        x = self.fc(x[-1, :, :])
        return x

数据准备

我们将使用随机生成的数据来训练模型。具体的实现如下：

def create_dataset(seq_length, num_samples, n_features):
    X = torch.randn(num_samples, seq_length, n_features)
    Y = torch.randn(num_samples, n_features)
    return X, Y

# 设置参数
seq_length = 10
num_samples = 100
n_features = 16
n_output = 16

X, Y = create_dataset(seq_length, num_samples, n_features)

模型训练

接下来，我们将定义训练过程并训练模型：

def train(model, X, Y, epochs=100, lr=0.001):
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()
        output = model(X.permute(1, 0, 2))  # 调整输入维度
        loss = criterion(output, Y)
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')

# 初始化模型并训练
model = TransformerModel(n_features, n_heads=4, n_layers=2, n_output=n_output)
train(model, X, Y)