transformer实现多变量多步时间序列预测的代码

以下是一个使用Transformer实现多变量多步时间序列预测的代码示例：

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

class MultiVariateTransformer(nn.Module):
    def __init__(self, input_dim, output_dim, n_head, n_layers, dropout):
        super(MultiVariateTransformer, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.n_head = n_head
        self.n_layers = n_layers
        self.dropout = dropout
        
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=input_dim, nhead=n_head, dropout=dropout)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=n_layers)
        self.decoder_layer = nn.TransformerDecoderLayer(d_model=output_dim, nhead=n_head, dropout=dropout)
        self.transformer_decoder = nn.TransformerDecoder(self.decoder_layer, num_layers=n_layers)
        self.fc = nn.Linear(input_dim, output_dim)
        
    def forward(self, input_seq):
        # input_seq: (seq_len, batch_size, input_dim)
        input_seq = input_seq.permute(1, 0, 2) # (batch_size, seq_len, input_dim)
        encoder_output = self.transformer_encoder(input_seq) # (batch_size, seq_len, input_dim)
        decoder_input = torch.zeros((input_seq.shape[0], self.output_dim)).to(input_seq.device) # (batch_size, output_dim)
        decoder_input = decoder_input.unsqueeze(1).repeat(1, input_seq.shape[1], 1) # (batch_size, seq_len, output_dim)
        for i in range(decoder_input.shape[1]):
            decoder_output = self.transformer_decoder(decoder_input[:, :i+1, :], encoder_output) # (batch_size, i+1, output_dim)
            decoder_input[:, i, :] = self.fc(decoder_output[:, i, :]) # (batch_size, output_dim)
        return decoder_input.permute(1, 0, 2) # (seq_len, batch_size, output_dim)

使用示例：

# 准备数据
train_input_seq = np.random.rand(100, 64, 5) # (seq_len, batch_size, input_dim)
train_target_seq = np.random.rand(10, 64, 2) # (seq_len, batch_size, output_dim)
train_input_seq = torch.tensor(train_input_seq).float()
train_target_seq = torch.tensor(train_target_seq).float()

# 初始化模型
model = MultiVariateTransformer(input_dim=5, output_dim=2, n_head=8, n_layers=6, dropout=0.1)

# 训练模型
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()
for epoch in range(100):
    optimizer.zero_grad()
    pred_target_seq = model(train_input_seq)
    loss = criterion(pred_target_seq[-10:, :, :], train_target_seq)
    loss.backward()
    optimizer.step()
    print("Epoch:", epoch, "loss:", loss.item())

# 预测结果
test_input_seq = np.random.rand(100, 1, 5)
test_input_seq = torch.tensor(test_input_seq).float()
with torch.no_grad():
    pred_target_seq = model(test_input_seq)
print(pred_target_seq.shape) # (seq_len, batch_size, output_dim)