使用pytorch搭建 transformer时间序列预测

Transformer模型可以用于时间序列预测，可以通过输入之前的时间步骤的值来预测未来的值。以下是使用PyTorch搭建Transformer模型进行时间序列预测的示例代码：

首先导入必要的库：

import torch
import torch.nn as nn
import torch.optim as optim
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_layers):
        super(TransformerModel, self).__init__()
        
        self.d_model = d_model
        self.nhead = nhead
        self.num_layers = num_layers
        
        # Encoder layers
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        
        # Decoder layers
        self.decoder_layer = nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead)
        self.decoder = nn.TransformerDecoder(self.decoder_layer, num_layers=num_layers)
        
        # Input and output embedding layers
        self.input_embedding = nn.Linear(input_dim, d_model)
        self.output_embedding = nn.Linear(d_model, output_dim)
        
    def forward(self, src, tgt):
        src = self.input_embedding(src.transpose(0, 1))
        tgt = self.input_embedding(tgt.transpose(0, 1))
        
        # Encoder
        memory = self.encoder(src)
        
        # Decoder
        output = self.decoder(tgt, memory)
        output = self.output_embedding(output.transpose(0, 1))
        
        return output

在此模型中，我们使用TransformerEncoder和TransformerDecoder来构建Transformer模型。输入和输出都通过一个线性层进行嵌入。在正向传递期间，我们首先将输入和输出嵌入传递给Encoder和Decoder，然后将Encoder输出传递给Decoder，最后将Decoder输出进行解嵌入以获得最终的预测输出。

接下来，我们定义模型的训练和测试函数：

def train(model, train_loader, optimizer, criterion, device):
    model.train()
    
    train_loss = 0
    
    for src, tgt in train_loader:
        src = src.to(device)
        tgt = tgt.to(device)
        
        optimizer.zero_grad()
        
        output = model(src, tgt)
        loss = criterion(output[:, -1, :], tgt[:, -1, :])
        loss.backward()
        
        optimizer.step()
        
        train_loss += loss.item()
        
    return train_loss / len(train_loader)


def test(model, test_loader, criterion, device):
    model.eval()
    
    test_loss = 0
    
    with torch.no_grad():
        for src, tgt in test_loader:
            src = src.to(device)
            tgt = tgt.to(device)

            output = model(src, tgt)
            loss = criterion(output[:, -1, :], tgt[:, -1, :])

            test_loss += loss.item()
        
    return test_loss / len(test_loader)

在训练函数中，我们首先将模型设置为训练模式，然后迭代数据集，并对每个数据点进行正向传递和反向传递。在测试函数中，我们使用相同的方式对数据集进行迭代，但是不进行反向传递。

最后，我们可以加载数据并训练模型：

# Load data
data = np.sin(np.arange(0, 1000, 0.1))
train_data = data[:8000]
test_data = data[8000:]

# Create train and test loaders
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=64, shuffle=False)

# Define hyperparameters
input_dim = 1
output_dim = 1
d_model = 64
nhead = 2
num_layers = 2
lr = 0.001
num_epochs = 50

# Initialize model, optimizer, and loss function
model = TransformerModel(input_dim, output_dim, d_model, nhead, num_layers)
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.MSELoss()

# Move model to device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Train model
train_losses = []
test_losses = []

for epoch in range(num_epochs):
    train_loss = train(model, train_loader, optimizer, criterion, device)
    test_loss = test(model, test_loader, criterion, device)
    
    print(f'Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')
    
    train_losses.append(train_loss)
    test_losses.append(test_loss)

# Plot loss curves
plt.plot(train_losses, label='Train Loss')
plt.plot(test_losses, label='Test Loss')
plt.legend()
plt.show()

在此代码中，我们首先生成了一个正弦波作为时间序列数据，然后创建了一个训练和测试数据加载器。接下来，我们定义了超参数，并初始化了模型、优化器和损失函数。最后，我们将模型移动到设备上并训练模型。我们还绘制了训练和测试损失曲线以进行可视化。