用transformer实现时间序列预测的代码

以下是一个用Transformer进行时间序列预测的Python代码示例。该示例使用PyTorch深度学习框架。

首先，导入所需的Python库和模块：

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

接下来，创建一个Transformer模型：

class TransformerModel(nn.Module):
    def __init__(self, input_dim, output_dim, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout):
        super(TransformerModel, self).__init__()
        
        # Encoder layer
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        
        # Decoder layer
        decoder_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
        
        # Linear layer
        self.linear = nn.Linear(d_model, output_dim)
        
        # Positional encoding
        self.pos_encoder = nn.Embedding(input_dim, d_model)
        self.pos_decoder = nn.Embedding(output_dim, d_model)
        
        # Dropout
        self.dropout = nn.Dropout(p=dropout)
        
        # D_model
        self.d_model = d_model
        
    def forward(self, src, tgt):
        # Add positional encoding to input and target
        src = self.pos_encoder(src) * np.sqrt(self.d_model)
        tgt = self.pos_decoder(tgt) * np.sqrt(self.d_model)
        
        # Apply dropout
        src = self.dropout(src)
        tgt = self.dropout(tgt)
        
        # Encoder output
        memory = self.encoder(src)
        
        # Decoder output
        output = self.decoder(tgt, memory)
        
        # Linear layer
        output = self.linear(output)
        
        return output

然后，准备数据并训练模型：

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Prepare data
train_data = torch.from_numpy(train_data).to(device)
train_target = torch.from_numpy(train_target).to(device)
test_data = torch.from_numpy(test_data).to(device)
test_target = torch.from_numpy(test_target).to(device)

# Model parameters
input_dim = 1
output_dim = 1
d_model = 256
nhead = 8
num_encoder_layers = 6
num_decoder_layers = 6
dim_feedforward = 512
dropout = 0.1

# Create model
model = TransformerModel(input_dim, output_dim, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout).to(device)

# Loss function
criterion = nn.MSELoss()

# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Train model
num_epochs = 100
for epoch in range(num_epochs):
    # Forward pass
    outputs = model(train_data, train_target[:-1])
    loss = criterion(outputs, train_target[1:])
    
    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # Print training progress
    if (epoch+1) % 10 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))

最后，使用训练好的模型进行预测并画图：

# Test model
model.eval()
with torch.no_grad():
    # Predict future values
    future = 100
    pred = model(test_data, test_target[-future:-1]).detach().cpu().numpy()
    
    # Plot results
    plt.figure(figsize=(12,5))
    plt.plot(test_target.cpu().numpy(), label='Observed')
    plt.plot(np.arange(len(test_target)-future, len(test_target)), pred, label='Predicted')
    plt.legend(loc='upper left')
    plt.show()