Merge pull request zhouhaoyi#34 from cookieminions/dev

Update features MS and fix custom data error
salvivona · Feb 24, 2021 · 2948c0e · 2948c0e
2 parents 8568425 + 9d74efd
commit 2948c0e
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 12 deletions.
diff --git a/data/data_loader.py b/data/data_loader.py
@@ -50,7 +50,7 @@ def __read_data__(self):
         border1 = border1s[self.set_type]
         border2 = border2s[self.set_type]
 
-        if self.features=='M':
+        if self.features=='M' or self.features=='MS':
             cols_data = df_raw.columns[1:]
             df_data = df_raw[cols_data]
         elif self.features=='S':
@@ -132,7 +132,7 @@ def __read_data__(self):
         border1 = border1s[self.set_type]
         border2 = border2s[self.set_type]
 
-        if self.features=='M':
+        if self.features=='M' or self.features=='MS':
             cols_data = df_raw.columns[1:]
             df_data = df_raw[cols_data]
         elif self.features=='S':
@@ -225,7 +225,7 @@ def __read_data__(self):
         border1 = border1s[self.set_type]
         border2 = border2s[self.set_type]
 
-        if self.features=='M':
+        if self.features=='M' or self.features=='MS':
             cols_data = df_raw.columns[1:]
             df_data = df_raw[cols_data]
         elif self.features=='S':

diff --git a/exp/exp_informer.py b/exp/exp_informer.py
@@ -77,6 +77,7 @@ def _get_data(self, flag):
             flag=flag,
             size=[args.seq_len, args.label_len, args.pred_len],
             features=args.features,
+            target=args.target,
             timeenc=timeenc,
             freq=args.freq
         )
@@ -116,7 +117,8 @@ def vali(self, vali_data, vali_loader, criterion):
                 outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
             else:
                 outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
-            batch_y = batch_y[:,-self.args.pred_len:,:].to(self.device)
+            f_dim = -1 if self.args.features=='MS' else 0
+            batch_y = batch_y[:,-self.args.pred_len:,f_dim:].to(self.device)
 
             pred = outputs.detach().cpu()
             true = batch_y.detach().cpu()
@@ -170,7 +172,8 @@ def train(self, setting):
                 else:
                     outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
 
-                batch_y = batch_y[:,-self.args.pred_len:,:].to(self.device)
+                f_dim = -1 if self.args.features=='MS' else 0
+                batch_y = batch_y[:,-self.args.pred_len:,f_dim:].to(self.device)
                 loss = criterion(outputs, batch_y)
                 train_loss.append(loss.item())
 
@@ -225,7 +228,8 @@ def test(self, setting):
                 outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)[0]
             else:
                 outputs = self.model(batch_x, batch_x_mark, dec_inp, batch_y_mark)
-            batch_y = batch_y[:,-self.args.pred_len:,:].to(self.device)
+            f_dim = -1 if self.args.features=='MS' else 0
+            batch_y = batch_y[:,-self.args.pred_len:,f_dim:].to(self.device)
 
             pred = outputs.detach().cpu().numpy()#.squeeze()
             true = batch_y.detach().cpu().numpy()#.squeeze()

diff --git a/main_informer.py b/main_informer.py
@@ -54,10 +54,10 @@
 args.use_gpu = True if torch.cuda.is_available() else False
 
 data_parser = {
-    'ETTh1':{'data':'ETTh1.csv','T':'OT','M':[7,7,7],'S':[1,1,1]},
-    'ETTh2':{'data':'ETTh2.csv','T':'OT','M':[7,7,7],'S':[1,1,1]},
-    'ETTm1':{'data':'ETTm1.csv','T':'OT','M':[7,7,7],'S':[1,1,1]},
-    'ETTm2':{'data':'ETTm2.csv','T':'OT','M':[7,7,7],'S':[1,1,1]},
+    'ETTh1':{'data':'ETTh1.csv','T':'OT','M':[7,7,7],'S':[1,1,1],'MS':[7,7,1]},
+    'ETTh2':{'data':'ETTh2.csv','T':'OT','M':[7,7,7],'S':[1,1,1],'MS':[7,7,1]},
+    'ETTm1':{'data':'ETTm1.csv','T':'OT','M':[7,7,7],'S':[1,1,1],'MS':[7,7,1]},
+    'ETTm2':{'data':'ETTm2.csv','T':'OT','M':[7,7,7],'S':[1,1,1],'MS':[7,7,1]},
 }
 if args.data in data_parser.keys():
     data_info = data_parser[args.data]

diff --git a/models/embed.py b/models/embed.py
@@ -86,7 +86,7 @@ class TimeFeatureEmbedding(nn.Module):
     def __init__(self, d_model, embed_type='timeF', freq='h'):
         super(TimeFeatureEmbedding, self).__init__()
 
-        freq_map = {'h':4, 't':5}
+        freq_map = {'h':4, 't':5, 'm':1, 'a':1, 'w':2, 'd':3, 'b':3}
         d_inp = freq_map[freq]
         self.embed = nn.Linear(d_inp, d_model)
 

diff --git a/utils/masking.py b/utils/masking.py
@@ -21,4 +21,4 @@ def __init__(self, B, H, L, index, scores, device="cpu"):
 
     @property
     def mask(self):
-        return self._mask
+        return self._mask