first commit

models/TSMixer.py

@@ -0,0 +1,54 @@
+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')