Spaces:
Sleeping
Sleeping
File size: 4,630 Bytes
e1ccef5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
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')
|