Fix LSTM benchmark to evaluate on test set

examples/mackey_glass/lstm_benchmark.py

@@ -164,10 +164,16 @@
         wandb.log({'fig_train': wandb.Image(fig)})
         plt.close()
 
+    # Initialize hidden states with the training data,
+    # now that weights are trained
+    lstm(train_data.to(device))
+
     # Testing
-    test_set_loader = DataLoader(train_set,
+    test_set_loader = DataLoader(test_set,
                                  batch_size=mg.testtime_pts,
                                  shuffle=False)
+
+
     lstm.mode = "autonomous"
     lstm.device = torch.device("cpu")
     lstm.to(torch.device("cpu"))
@@ -210,9 +216,8 @@
       f"on tau {args.tau}")
 
 # With the default params, repeat 30, tau=17
-# sMAPE score = 15.156239883579927,
+# sMAPE score = 14.347067906362048,
 # connection_sparsity = 0.0,
-# activation_sparsity = 0.45951777777777786,
-# synop_macs = 14534.032622222225,
-# synop_dense = 14552.413333333332,
-# on time series id 0
+# activation_sparsity = 0.46036666666666665,
+# synop_macs = 14541.585333333329,
+# synop_dense = 14560.0,

examples/mackey_glass/lstm_model.py

@@ -111,16 +111,22 @@ def forward(self, batch):
             3D tesnor: output data of shape [num_steps, output_dim]
         """
 
-        # Reset register buffers
-        for i in range(self.n_layers):
-            getattr(self, f'h{i}')[:] = 0.
-        for i in range(self.n_layers):
-            getattr(self, f'c{i}')[:] = 0.
-        self.inp[:] = 0.
-
-        self.h = [getattr(self, f'h{i}') for i in range(self.n_layers)]
-        self.c = [getattr(self, f'c{i}') for i in range(self.n_layers)]
-
+        if self.mode != 'autonomous':
+            # Reset register buffers
+            for i in range(self.n_layers):
+                getattr(self, f'h{i}')[:] = 0.
+            for i in range(self.n_layers):
+                getattr(self, f'c{i}')[:] = 0.
+
+            self.h = [getattr(self, f'h{i}') for i in range(self.n_layers)]
+            self.c = [getattr(self, f'c{i}') for i in range(self.n_layers)]
+            self.inp[:] = 0.
+        else:
+            # The warmed-up LSTM states and the buffer at the end of training are used during
+            # autonomous replay (for an explanation, see section 3.1 in https://doi.org/10.1016/j.mlwa.2022.100300)
+            self.h = [self.h[i].to(batch.device) for i in range(self.n_layers)]
+            self.c = [self.c[i].to(batch.device) for i in range(self.n_layers)]
+
         predictions = []
         for i, sample in enumerate(batch):