### Informer 时间序列预测示例代码
以下是基于 PyTorch 的 Informer 模型实现时间序列预测的一个简单示例。此代码展示了如何构建和训练一个基本的 Informer 模型来完成多步预测。
```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import numpy as np
class TimeSeriesDataset(Dataset):
def __init__(self, data, seq_len=96, pred_len=24):
self.seq_len = seq_len
self.pred_len = pred_len
self.data = data
def __len__(self):
return len(self.data) - self.seq_len - self.pred_len + 1
def __getitem__(self, idx):
x = self.data[idx:idx+self.seq_len]
y = self.data[idx+self.seq_len:idx+self.seq_len+self.pred_len]
return torch.tensor(x, dtype=torch.float), torch.tensor(y, dtype=torch.float)
class ProbSparseAttention(nn.Module):
"""概率稀疏注意力机制"""
def forward(self, queries, keys, values):
B, L_Q, D = queries.shape
_, L_K, _ = keys.shape
scores = torch.matmul(queries, keys.transpose(-2, -1)) / (D ** 0.5)
top_k_scores, indices = torch.topk(scores, k=int(L_K * 0.5), dim=-1)
sparse_mask = torch.zeros_like(scores).scatter_(-1, indices, top_k_scores)
attention_weights = torch.softmax(sparse_mask, dim=-1)
context_vector = torch.matmul(attention_weights, values)
return context_vector
class EncoderLayer(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.attention = ProbSparseAttention()
self.linear_projection = nn.Linear(d_model, d_model)
def forward(self, x):
attn_output = self.attention(x, x, x)
output = self.linear_projection(attn_output)
return output
class DecoderLayer(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.self_attention = ProbSparseAttention()
self.cross_attention = ProbSparseAttention()
self.linear_projection = nn.Linear(d_model, d_model)
def forward(self, x, enc_output):
self_attn_output = self.self_attention(x, x, x)
cross_attn_output = self.cross_attention(self_attn_output, enc_output, enc_output)
output = self.linear_projection(cross_attn_output)
return output
class InformerModel(nn.Module):
def __init__(self, input_size, hidden_size, output_size, num_layers, n_heads):
super(InformerModel, self).__init__()
self.encoder_embedding = nn.Linear(input_size, hidden_size)
self.decoder_embedding = nn.Linear(input_size, hidden_size)
self.encoder_layers = nn.ModuleList([EncoderLayer(hidden_size, n_heads) for _ in range(num_layers)])
self.decoder_layers = nn.ModuleList([DecoderLayer(hidden_size, n_heads) for _ in range(num_layers)])
self.final_linear = nn.Linear(hidden_size, output_size)
def forward(self, src, tgt):
encoded_src = self.encoder_embedding(src.permute(0, 2, 1)).permute(0, 2, 1)
decoded_tgt = self.decoder_embedding(tgt.permute(0, 2, 1)).permute(0, 2, 1)
encoder_outputs = encoded_src
for layer in self.encoder_layers:
encoder_outputs = layer(encoder_outputs)
decoder_outputs = decoded_tgt
for layer in self.decoder_layers:
decoder_outputs = layer(decoder_outputs, encoder_outputs)
final_output = self.final_linear(decoder_outputs)
return final_output.squeeze()
def train(model, dataloader, criterion, optimizer, device='cpu'):
model.train()
total_loss = 0
for batch_x, batch_y in dataloader:
batch_x, batch_y = batch_x.to(device), batch_y.to(device)
outputs = model(batch_x.unsqueeze(-1), batch_y[:, :-1].unsqueeze(-1))
loss = criterion(outputs, batch_y[:, 1:].unsqueeze(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
return avg_loss
# 数据准备
data = np.sin(np.linspace(0, 100, 1000)) # 使用正弦波作为模拟数据
dataset = TimeSeriesDataset(data=data, seq_len=96, pred_len=24)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
# 初始化模型、损失函数和优化器
model = InformerModel(input_size=1, hidden_size=64, output_size=1, num_layers=2, n_heads=8)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
for epoch in range(10): # 训练10轮
loss = train(model, dataloader, criterion, optimizer)
print(f'Epoch {epoch+1}, Loss: {loss:.4f}')
```
上述代码实现了 Informer 模型的核心部分，包括编码层、解码层以及概率稀疏注意力机制[^1]。通过调整超参数（如 `seq_len` 和 `pred_len`），可以适应不同的时间序列长度需求。
#### 注意事项
- 上述代码中的 `ProbSparseAttention` 是简化版的概率稀疏注意力模块，实际应用中可能需要更复杂的实现。
- 正弦波被用作测试数据集；真实场景下应替换为具体业务的数据源。
文章介绍了如何使用Informer模型进行时间序列预测，该模型基于Transformer，适用于时间序列任务。通过展示从数据准备到模型训练再到预测的流程，说明了Informer在预测未来趋势上的有效性，并强调了其在实际应用中的价值。
订阅专栏 解锁全文
扫一扫
7334
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
