写一下transformer时间序列预测的代码

Transformer是一种强大的神经网络模型，常用于处理序列数据，包括时间序列预测。下面是一个简单的Transformer时间序列预测的代码示例：

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

# 定义Transformer模型
class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim, num_layers, num_heads):
        super(Transformer, self).__init__()
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(input_dim, num_heads, hidden_dim),
            num_layers
        )
        self.decoder = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x[:, -1, :])  # 只使用最后一个时间步的输出进行预测
        return x

# 准备数据
input_dim = 10  # 输入维度
output_dim = 1  # 输出维度
hidden_dim = 32  # 隐藏层维度
num_layers = 2  # Transformer编码器层数
num_heads = 4  # 注意力头数

# 创建模型和优化器
model = Transformer(input_dim, output_dim, hidden_dim, num_layers, num_heads)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# 训练模型
num_epochs = 100
for epoch in range(num_epochs):
    optimizer.zero_grad()
    inputs = torch.randn(100, input_dim)  # 随机生成输入数据
    targets = torch.randn(100, output_dim)  # 随机生成目标数据
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

# 使用模型进行预测
test_inputs = torch.randn(10, input_dim)  # 随机生成测试输入数据
with torch.no_grad():
    test_outputs = model(test_inputs)
print('Test Outputs:', test_outputs)

这段代码实现了一个简单的Transformer模型，用于时间序列预测。首先定义了一个Transformer类，其中包含了一个Transformer编码器和一个线性层用于预测。然后准备数据，创建模型和优化器。接下来进行模型训练，使用随机生成的输入数据和目标数据进行训练，并输出每个epoch的损失值。最后使用训练好的模型进行预测，输出测试数据的预测结果。