pytorch Transformer交通预测

交通预测可以看作是时间序列预测问题，可以使用Transformer模型进行建模。具体来说，可以将历史时间窗口内的交通数据作为输入序列，预测未来一段时间内的交通情况。

以下是一个PyTorch中使用Transformer模型进行交通预测的示例代码：

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

# 定义Transformer模型
class TransformerModel(nn.Module):
    def __init__(self, input_size, output_size, d_model, nhead, num_layers):
        super(TransformerModel, self).__init__()
        self.encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead)
        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
        self.fc = nn.Linear(d_model, output_size)

    def forward(self, src):
        src = src.permute(1, 0, 2)  # 调整输入数据的维度
        output = self.transformer_encoder(src)
        output = output.permute(1, 0, 2)  # 调整输出数据的维度
        output = self.fc(output[:, -1, :])  # 取最后一个时间步的输出作为预测结果
        return output

# 定义数据集类
class TrafficDataset(torch.utils.data.Dataset):
    def __init__(self, data, seq_len, target_len):
        self.data = data
        self.seq_len = seq_len
        self.target_len = target_len

    def __len__(self):
        return len(self.data) - self.seq_len - self.target_len

    def __getitem__(self, idx):
        x = self.data[idx:idx+self.seq_len]
        y = self.data[idx+self.seq_len:idx+self.seq_len+self.target_len]
        return x, y

# 准备数据
data = np.random.rand(100, 2)  # 假设有100组交通数据，每组数据包含2个特征
seq_len = 10  # 定义输入序列的长度
target_len = 5  # 定义目标序列的长度
train_data = TrafficDataset(data[:80], seq_len, target_len)
test_data = TrafficDataset(data[80:], seq_len, target_len)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=16)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=16)

# 定义模型和优化器
model = TransformerModel(input_size=2, output_size=target_len, d_model=64, nhead=4, num_layers=2)
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    for x, y in train_loader:
        optimizer.zero_grad()
        output = model(x.float())
        loss = nn.MSELoss()(output, y.float())
        loss.backward()
        optimizer.step()
    print('Epoch {}, Loss: {:.4f}'.format(epoch, loss.item()))

# 测试模型
with torch.no_grad():
    for x, y in test_loader:
        output = model(x.float())
        print('Predicted: ', output)
        print('Ground truth: ', y)

在上面的示例中，我们使用了PyTorch中自带的`TransformerEncoder`和`TransformerEncoderLayer`模块来构建Transformer模型，并且使用MSE作为损失函数进行训练。同时，我们还通过自定义数据集类和数据加载器来准备训练和测试数据。