| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
|
|
| from layers.embed import Time2Vec |
|
|
|
|
| class MLP(nn.Module): |
| """ |
| Just your everyday neural net |
| """ |
|
|
| def __init__(self, config): |
| super(MLP, self).__init__() |
| self.seq_len = config.seq_len |
| self.pred_len = config.pred_len |
| assert config.pred_len == 1 |
|
|
| self.e_layers = config.e_layers |
| assert config.e_layers >= 1 |
|
|
| self.enc_in = config.enc_in |
| self.d_model = config.d_model |
| self.c_out = config.c_out |
|
|
| |
| self.app_time_emb = config.t_embed is not None |
| if self.app_time_emb: |
| if config.t_embed != "time2vec_app": |
| raise Exception( |
| "The only options for t_embed with mlp are null and time2vec_app" |
| ) |
| elif not (config.emb_t2v_app_dim > 0): |
| raise Exception("Need to specify a valid emb_t2v_app_dim") |
| self.enc_in += config.emb_t2v_app_dim |
| self.temporal_embedding = Time2Vec( |
| time_emb_dim=config.emb_t2v_app_dim, freq=config.freq |
| ) |
|
|
| flattened_enc_in = self.seq_len * self.enc_in |
|
|
| if self.e_layers == 1: |
| layers = [nn.Linear(flattened_enc_in, self.c_out)] |
| else: |
| layers = [nn.Linear(flattened_enc_in, self.d_model), nn.GELU()] |
| for _ in range(self.e_layers - 2): |
| layers.append(nn.Dropout(config.dropout)) |
| layers.append(nn.Linear(self.d_model, self.d_model)) |
| layers.append(nn.GELU()) |
|
|
| layers.append(nn.Linear(self.d_model, self.c_out)) |
|
|
| self.model = nn.Sequential(*layers) |
|
|
| def forward(self, x, x_mark, *args): |
| |
| if self.app_time_emb: |
| time_emb = self.temporal_embedding(x_mark) |
| x = torch.concat([x, time_emb], dim=-1) |
|
|
| x_flat = x.reshape(x.shape[0], 1, -1) |
| return self.model(x_flat) |
|
|
