| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import numpy as np
|
| from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer
|
| from layers.SelfAttention_Family import FullAttention, AttentionLayer
|
| from layers.Embed import PatchEmbedding
|
| from collections import Counter
|
| from layers.SharedWavMoE import WavMoE
|
| from layers.RevIN import RevIN
|
| import torch.fft
|
| from layers.Embed import DataEmbedding
|
|
|
| class FlattenHead(nn.Module):
|
| def __init__(self, n_vars, nf, target_window, head_dropout=0):
|
| super().__init__()
|
| self.n_vars = n_vars
|
|
|
| self.linear = nn.Linear(nf, target_window)
|
| self.dropout = nn.Dropout(head_dropout)
|
|
|
| def forward(self, x):
|
|
|
|
|
| x = self.linear(x)
|
| x = self.dropout(x)
|
| return x
|
|
|
|
|
|
|
| class Model(nn.Module):
|
| """
|
| """
|
|
|
| def __init__(self, configs):
|
| super(Model, self).__init__()
|
| self.task_name = configs.task_name
|
| self.seq_len = configs.seq_len
|
| self.patch_len = configs.input_token_len
|
| self.stride = self.patch_len
|
| self.pred_len = configs.test_pred_len
|
| self.test_seq_len = configs.test_seq_len
|
|
|
| self.output_attention = configs.output_attention
|
| self.padding = configs.padding
|
|
|
| self.hidden_size = configs.hidden_size
|
| self.intermediate_size = configs.intermediate_size
|
| self.top_k = configs.top_k
|
| self.shared_experts = configs.shared_experts
|
| self.wavelet = configs.wavelet
|
| self.level = configs.shared_experts
|
| self.proj_wight = configs.proj_wight
|
|
|
| self.patch_embedding = PatchEmbedding(
|
| configs.d_model, self.patch_len, self.stride, self.padding, configs.dropout)
|
|
|
| self.data_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
|
| configs.dropout)
|
|
|
| self.revin_layer = RevIN(configs.enc_in)
|
| self.encoder_patch = Encoder(
|
| [
|
| EncoderLayer(
|
| AttentionLayer(
|
| FullAttention(False, configs.factor,
|
| attention_dropout=configs.dropout,
|
| output_attention=configs.output_attention),
|
| configs.d_model, configs.n_heads),
|
| configs.d_model,
|
| configs.d_ff,
|
| dropout=configs.dropout,
|
| activation=configs.activation
|
| ) for l in range(configs.e_layers)
|
| ],
|
| norm_layer=torch.nn.LayerNorm(configs.d_model)
|
| )
|
| self.encoder_time = Encoder(
|
| [
|
| EncoderLayer(
|
| AttentionLayer(
|
| FullAttention(False, configs.factor,
|
| attention_dropout=configs.dropout,
|
| output_attention=configs.output_attention),
|
| configs.d_model, configs.n_heads),
|
| configs.d_model,
|
| configs.d_ff,
|
| dropout=configs.dropout,
|
| activation=configs.activation
|
| ) for l in range(configs.e_layers)
|
| ],
|
| norm_layer=torch.nn.LayerNorm(configs.d_model)
|
| )
|
| self.head_nf = configs.d_model * \
|
| int((configs.seq_len - self.patch_len) / self.stride + 1)
|
| self.projection = nn.Linear(self.head_nf, int(configs.seq_len*self.proj_wight), bias=True)
|
|
|
| self.data_projection = nn.Linear(configs.d_model, configs.enc_in, bias=True)
|
| self.wavmoe = WavMoE(configs)
|
| self.head = FlattenHead(configs.enc_in, nf= int(configs.seq_len*self.proj_wight), target_window= self.seq_len,
|
| head_dropout=configs.dropout)
|
| self.gelu = nn.GELU()
|
|
|
|
|
|
|
| def main(self, x_enc, x_mark_enc, x_dec, x_mark_dec):
|
|
|
| x_revin = self.revin_layer(x_enc, 'norm').permute(0, 2, 1)
|
|
|
| B, D, S = x_revin.shape
|
|
|
|
|
| x_inver=self.data_embedding(x_revin.permute(0, 2, 1), x_mark_enc)
|
| nav_out, attn_w = self.encoder_time(x_inver, attn_mask=None)
|
|
|
| nav_out = self.data_projection(nav_out)
|
|
|
|
|
|
|
|
|
|
|
| x_pe, n_vars = self.patch_embedding(x_revin+nav_out.permute(0, 2, 1))
|
|
|
| enc_out, attn = self.encoder_patch(x_pe)
|
| dec_out = enc_out.reshape(B, D, -1)
|
|
|
| act_val = self.projection(dec_out)
|
|
|
|
|
|
|
| moe_out, router_logits = self.wavmoe(act_val + nav_out.permute(0, 2, 1))
|
|
|
| head_out = self.head(moe_out)
|
|
|
|
|
| x_out = self.revin_layer(head_out.permute(0, 2, 1), 'denorm')
|
|
|
| return x_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 == 'forecast':
|
| dec_out = self.main(x_enc, x_mark_enc, x_dec, x_mark_dec)
|
| return dec_out[:, -self.test_seq_len :, :]
|
| if self.task_name == 'anomaly_detection':
|
| dec_out = self.main(x_enc, x_mark_enc, x_dec, x_mark_dec)
|
| return dec_out
|
| return None
|
|
|
|
|
