transformer时间序列预测的代码

下面是一个用Transformer模型进行时间序列预测的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# 定义Transformer模型
class Transformer(nn.Module):
    def __init__(self, input_size, output_size, hidden_size, num_layers, dropout):
        super(Transformer, self).__init__()

        self.input_size = input_size
        self.output_size = output_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.dropout = dropout

        self.encoder_layer = nn.TransformerEncoderLayer(d_model=input_size, nhead=1, dropout=dropout)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        self.decoder_layer = nn.TransformerDecoderLayer(d_model=output_size, nhead=1, dropout=dropout)
        self.decoder = nn.TransformerDecoder(self.decoder_layer, num_layers=num_layers)

        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, y):
        # x: [seq_len, batch_size, input_size]
        # y: [seq_len, batch_size, output_size]

        # 编码器输入为x，解码器输入为y[:-1]
        x = self.encoder(x)
        y = self.decoder(y[:-1], x)

        # 解码器输出经过全连接层
        y = self.fc(y)

        return y

# 生成数据
def generate_data(data_size, seq_len):
    x = np.zeros((seq_len, data_size))
    y = np.zeros((seq_len, data_size))

    for i in range(seq_len):
        x[i] = np.sin(np.linspace(i*0.1, (i+1)*0.1, data_size))
        y[i] = np.sin(np.linspace((i+1)*0.1, (i+2)*0.1, data_size))

    return x, y

# 定义超参数
input_size = 10
output_size = 10
hidden_size = 128
num_layers = 2
dropout = 0.1
learning_rate = 0.001
batch_size = 32
seq_len = 100

# 生成数据
x, y = generate_data(input_size, seq_len)
x = torch.tensor(x, dtype=torch.float32).unsqueeze(1)  # [seq_len, 1, input_size]
y = torch.tensor(y, dtype=torch.float32).unsqueeze(1)  # [seq_len, 1, output_size]

# 定义模型、优化器和损失函数
model = Transformer(input_size, output_size, hidden_size, num_layers, dropout)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.MSELoss()

# 训练模型
for epoch in range(100):
    epoch_loss = 0

    for i in range(0, seq_len-batch_size, batch_size):
        batch_x = x[i:i+batch_size]  # [batch_size, 1, input_size]
        batch_y = y[i:i+batch_size]  # [batch_size, 1, output_size]
        pred_y = model(batch_x, batch_y)  # [seq_len-1, batch_size, output_size]

        loss = criterion(pred_y.view(-1, output_size), batch_y[1:].view(-1, output_size))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()

    print('Epoch %d loss: %.3f' % (epoch+1, epoch_loss))

# 预测未来数据
future_x = x[-1].unsqueeze(0)  # [1, 1, input_size]
future_y = model(future_x, torch.zeros(1, 1, output_size))  # [seq_len-1, 1, output_size]
future_y = future_y.squeeze(1).detach().numpy()  # [seq_len-1, output_size]

这个代码用了一个简单的sin函数生成的数据集，每个样本有10个特征，共生成100个样本。模型用了两个Transformer层，每个层都有一个Encoder和一个Decoder。在训练的时候，每次给模型输入一个batch的数据，输出的结果和实际的结果计算MSE损失，然后进行反向传播和优化。最后，用训练好的模型预测接下来的数据。