67 lines
2.1 KiB
Python
67 lines
2.1 KiB
Python
"""
|
|
SeasonPatch = PatchTST (CI) + ChannelGraphMixer + Linear prediction head
|
|
Adapted for Time-Series-Library-main style
|
|
"""
|
|
import torch
|
|
import torch.nn as nn
|
|
from layers.TSTEncoder import TSTiEncoder
|
|
from layers.GraphMixer import HierarchicalGraphMixer
|
|
|
|
class SeasonPatch(nn.Module):
|
|
def __init__(self,
|
|
c_in: int,
|
|
seq_len: int,
|
|
pred_len: int,
|
|
patch_len: int,
|
|
stride: int,
|
|
k_graph: int = 8,
|
|
d_model: int = 128,
|
|
n_layers: int = 3,
|
|
n_heads: int = 16):
|
|
super().__init__()
|
|
|
|
# Store patch parameters
|
|
self.patch_len = patch_len
|
|
self.stride = stride
|
|
|
|
# Calculate patch number
|
|
patch_num = (seq_len - patch_len) // stride + 1
|
|
|
|
# PatchTST encoder (channel independent)
|
|
self.encoder = TSTiEncoder(
|
|
c_in=c_in,
|
|
patch_num=patch_num,
|
|
patch_len=patch_len,
|
|
d_model=d_model,
|
|
n_layers=n_layers,
|
|
n_heads=n_heads
|
|
)
|
|
|
|
# Cross-channel mixer
|
|
self.mixer = HierarchicalGraphMixer(c_in, dim=d_model, k=k_graph)
|
|
|
|
# Prediction head
|
|
self.head = nn.Linear(patch_num * d_model, pred_len)
|
|
|
|
def forward(self, x):
|
|
# x: [B, L, C]
|
|
x = x.permute(0, 2, 1) # → [B, C, L]
|
|
|
|
# Patch the input
|
|
x_patch = x.unfold(-1, self.patch_len, self.stride) # [B, C, patch_num, patch_len]
|
|
|
|
# Encode patches
|
|
z = self.encoder(x_patch) # [B, C, d_model, patch_num]
|
|
|
|
# z: [B, C, d_model, patch_num] → [B, C, patch_num, d_model]
|
|
B, C, D, N = z.shape
|
|
z = z.permute(0, 1, 3, 2) # [B, C, patch_num, d_model]
|
|
|
|
# Cross-channel mixing
|
|
z_mix = self.mixer(z) # [B, C, patch_num, d_model]
|
|
|
|
# Flatten and predict
|
|
z_mix = z_mix.view(B, C, N * D) # [B, C, patch_num * d_model]
|
|
y_pred = self.head(z_mix) # [B, C, pred_len]
|
|
|
|
return y_pred |