Spaces:
Sleeping
Sleeping
File size: 1,814 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 |
import torch.nn as nn
class ResBlock(nn.Module):
def __init__(self, configs):
super(ResBlock, self).__init__()
self.temporal = nn.Sequential(
nn.Linear(configs.seq_len, configs.d_model),
nn.ReLU(),
nn.Linear(configs.d_model, configs.seq_len),
nn.Dropout(configs.dropout)
)
self.channel = nn.Sequential(
nn.Linear(configs.enc_in, configs.d_model),
nn.ReLU(),
nn.Linear(configs.d_model, configs.enc_in),
nn.Dropout(configs.dropout)
)
def forward(self, x):
# x: [B, L, D]
x = x + self.temporal(x.transpose(1, 2)).transpose(1, 2)
x = x + self.channel(x)
return x
class Model(nn.Module):
def __init__(self, configs):
super(Model, self).__init__()
self.task_name = configs.task_name
self.layer = configs.e_layers
self.model = nn.ModuleList([ResBlock(configs)
for _ in range(configs.e_layers)])
self.pred_len = configs.pred_len
self.projection = nn.Linear(configs.seq_len, configs.pred_len)
def forecast(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):
# x: [B, L, D]
for i in range(self.layer):
x_enc = self.model[i](x_enc)
enc_out = self.projection(x_enc.transpose(1, 2)).transpose(1, 2)
return enc_out
def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec, mask=None):
if self.task_name == 'long_term_forecast' or self.task_name == 'short_term_forecast':
dec_out = self.forecast(x_enc, x_mark_enc, x_dec, x_mark_dec)
return dec_out[:, -self.pred_len:, :] # [B, L, D]
else:
raise ValueError('Only forecast tasks implemented yet')
|