119 lines
4.5 KiB
Python
119 lines
4.5 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
|
|
|
|
class Model(nn.Module):
|
|
"""
|
|
Paper link: https://arxiv.org/pdf/2311.06184.pdf
|
|
"""
|
|
|
|
def __init__(self, configs):
|
|
super(Model, self).__init__()
|
|
self.task_name = configs.task_name
|
|
if self.task_name == 'classification' or self.task_name == 'anomaly_detection' or self.task_name == 'imputation':
|
|
self.pred_len = configs.seq_len
|
|
else:
|
|
self.pred_len = configs.pred_len
|
|
self.embed_size = 128 # embed_size
|
|
self.hidden_size = 256 # hidden_size
|
|
self.pred_len = configs.pred_len
|
|
self.feature_size = configs.enc_in # channels
|
|
self.seq_len = configs.seq_len
|
|
self.channel_independence = configs.channel_independence
|
|
self.sparsity_threshold = 0.01
|
|
self.scale = 0.02
|
|
self.embeddings = nn.Parameter(torch.randn(1, self.embed_size))
|
|
self.r1 = nn.Parameter(self.scale * torch.randn(self.embed_size, self.embed_size))
|
|
self.i1 = nn.Parameter(self.scale * torch.randn(self.embed_size, self.embed_size))
|
|
self.rb1 = nn.Parameter(self.scale * torch.randn(self.embed_size))
|
|
self.ib1 = nn.Parameter(self.scale * torch.randn(self.embed_size))
|
|
self.r2 = nn.Parameter(self.scale * torch.randn(self.embed_size, self.embed_size))
|
|
self.i2 = nn.Parameter(self.scale * torch.randn(self.embed_size, self.embed_size))
|
|
self.rb2 = nn.Parameter(self.scale * torch.randn(self.embed_size))
|
|
self.ib2 = nn.Parameter(self.scale * torch.randn(self.embed_size))
|
|
|
|
self.fc = nn.Sequential(
|
|
nn.Linear(self.seq_len * self.embed_size, self.hidden_size),
|
|
nn.LeakyReLU(),
|
|
nn.Linear(self.hidden_size, self.pred_len)
|
|
)
|
|
|
|
# dimension extension
|
|
def tokenEmb(self, x):
|
|
# x: [Batch, Input length, Channel]
|
|
x = x.permute(0, 2, 1)
|
|
x = x.unsqueeze(3)
|
|
# N*T*1 x 1*D = N*T*D
|
|
y = self.embeddings
|
|
return x * y
|
|
|
|
# frequency temporal learner
|
|
def MLP_temporal(self, x, B, N, L):
|
|
# [B, N, T, D]
|
|
x = torch.fft.rfft(x, dim=2, norm='ortho') # FFT on L dimension
|
|
y = self.FreMLP(B, N, L, x, self.r2, self.i2, self.rb2, self.ib2)
|
|
x = torch.fft.irfft(y, n=self.seq_len, dim=2, norm="ortho")
|
|
return x
|
|
|
|
# frequency channel learner
|
|
def MLP_channel(self, x, B, N, L):
|
|
# [B, N, T, D]
|
|
x = x.permute(0, 2, 1, 3)
|
|
# [B, T, N, D]
|
|
x = torch.fft.rfft(x, dim=2, norm='ortho') # FFT on N dimension
|
|
y = self.FreMLP(B, L, N, x, self.r1, self.i1, self.rb1, self.ib1)
|
|
x = torch.fft.irfft(y, n=self.feature_size, dim=2, norm="ortho")
|
|
x = x.permute(0, 2, 1, 3)
|
|
# [B, N, T, D]
|
|
return x
|
|
|
|
# frequency-domain MLPs
|
|
# dimension: FFT along the dimension, r: the real part of weights, i: the imaginary part of weights
|
|
# rb: the real part of bias, ib: the imaginary part of bias
|
|
def FreMLP(self, B, nd, dimension, x, r, i, rb, ib):
|
|
o1_real = torch.zeros([B, nd, dimension // 2 + 1, self.embed_size],
|
|
device=x.device)
|
|
o1_imag = torch.zeros([B, nd, dimension // 2 + 1, self.embed_size],
|
|
device=x.device)
|
|
|
|
o1_real = F.relu(
|
|
torch.einsum('bijd,dd->bijd', x.real, r) - \
|
|
torch.einsum('bijd,dd->bijd', x.imag, i) + \
|
|
rb
|
|
)
|
|
|
|
o1_imag = F.relu(
|
|
torch.einsum('bijd,dd->bijd', x.imag, r) + \
|
|
torch.einsum('bijd,dd->bijd', x.real, i) + \
|
|
ib
|
|
)
|
|
|
|
y = torch.stack([o1_real, o1_imag], dim=-1)
|
|
y = F.softshrink(y, lambd=self.sparsity_threshold)
|
|
y = torch.view_as_complex(y)
|
|
return y
|
|
|
|
def forecast(self, x_enc):
|
|
# x: [Batch, Input length, Channel]
|
|
B, T, N = x_enc.shape
|
|
# embedding x: [B, N, T, D]
|
|
x = self.tokenEmb(x_enc)
|
|
bias = x
|
|
# [B, N, T, D]
|
|
if self.channel_independence == '0':
|
|
x = self.MLP_channel(x, B, N, T)
|
|
# [B, N, T, D]
|
|
x = self.MLP_temporal(x, B, N, T)
|
|
x = x + bias
|
|
x = self.fc(x.reshape(B, N, -1)).permute(0, 2, 1)
|
|
return x
|
|
|
|
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
|
|
if self.task_name == 'long_term_forecast' or self.task_name == 'short_term_forecast':
|
|
dec_out = self.forecast(x_enc)
|
|
return dec_out[:, -self.pred_len:, :] # [B, L, D]
|
|
else:
|
|
raise ValueError('Only forecast tasks implemented yet')
|