Speed up testing procedure (#336)

* test: reduce the num of training epochs and models' sizes to speed up testing;

tests/classification/brits.py

@@ -44,7 +44,7 @@ class TestBRITS(unittest.TestCase):
         DATA["n_steps"],
         DATA["n_features"],
         n_classes=DATA["n_classes"],
-        rnn_hidden_size=256,
+        rnn_hidden_size=32,
         epochs=EPOCHS,
         saving_path=saving_path,
         model_saving_strategy="better",

tests/classification/grud.py

@@ -44,7 +44,7 @@ class TestGRUD(unittest.TestCase):
         DATA["n_steps"],
         DATA["n_features"],
         n_classes=DATA["n_classes"],
-        rnn_hidden_size=256,
+        rnn_hidden_size=32,
         epochs=EPOCHS,
         saving_path=saving_path,
         optimizer=optimizer,

tests/classification/raindrop.py

@@ -42,7 +42,7 @@ class TestRaindrop(unittest.TestCase):
         DATA["n_classes"],
         n_layers=2,
         d_model=DATA["n_features"] * 4,
-        d_ffn=256,
+        d_ffn=32,
         n_heads=2,
         dropout=0.3,
         d_static=0,

tests/clustering/crli.py

@@ -50,7 +50,7 @@ class TestCRLI(unittest.TestCase):
         n_features=DATA["n_features"],
         n_clusters=DATA["n_classes"],
         n_generator_layers=2,
-        rnn_hidden_size=128,
+        rnn_hidden_size=32,
         rnn_cell_type="GRU",
         epochs=EPOCHS,
         saving_path=saving_path,

tests/clustering/vader.py

@@ -49,7 +49,7 @@ class TestVaDER(unittest.TestCase):
         n_steps=DATA["n_steps"],
         n_features=DATA["n_features"],
         n_clusters=DATA["n_classes"],
-        rnn_hidden_size=64,
+        rnn_hidden_size=32,
         d_mu_stddev=5,
         pretrain_epochs=20,
         epochs=EPOCHS,

tests/global_test_config.py

@@ -51,7 +51,7 @@
 RESULT_SAVING_DIR_FOR_FORECASTING = os.path.join(RESULT_SAVING_DIR, "forecasting")
 
 # set the number of epochs for all model training
-EPOCHS = 5
+EPOCHS = 2
 
 # set DEVICES to None if no cuda device is available, to avoid initialization failed while importing test classes
 n_cuda_devices = torch.cuda.device_count()

tests/imputation/autoformer.py

@@ -47,8 +47,8 @@ class TestAutoformer(unittest.TestCase):
         DATA["n_features"],
         n_layers=2,
         n_heads=2,
-        d_model=128,
-        d_ffn=256,
+        d_model=32,
+        d_ffn=32,
         factor=3,
         moving_avg_window_size=3,
         dropout=0,

tests/imputation/brits.py

@@ -45,7 +45,7 @@ class TestBRITS(unittest.TestCase):
     brits = BRITS(
         DATA["n_steps"],
         DATA["n_features"],
-        256,
+        32,
         epochs=EPOCHS,
         saving_path=saving_path,
         optimizer=optimizer,

tests/imputation/crossformer.py

@@ -47,8 +47,8 @@ class TestCrossformer(unittest.TestCase):
         DATA["n_features"],
         n_layers=2,
         n_heads=2,
-        d_model=128,
-        d_ffn=256,
+        d_model=32,
+        d_ffn=32,
         factor=10,
         seg_len=12,
         win_size=2,

tests/imputation/csdi.py

@@ -49,7 +49,7 @@ class TestCSDI(unittest.TestCase):
         d_time_embedding=32,
         d_feature_embedding=3,
         d_diffusion_embedding=32,
-        n_diffusion_steps=10,
+        n_diffusion_steps=5,
         n_heads=1,
         epochs=EPOCHS,
         saving_path=saving_path,

tests/imputation/etsformer.py

@@ -48,8 +48,8 @@ class TestETSformer(unittest.TestCase):
         n_e_layers=2,
         n_d_layers=2,
         n_heads=2,
-        d_model=128,
-        d_ffn=256,
+        d_model=32,
+        d_ffn=32,
         top_k=3,
         dropout=0,
         epochs=EPOCHS,

tests/imputation/fedformer.py

@@ -47,8 +47,8 @@ class TestFEDformer(unittest.TestCase):
         DATA["n_features"],
         n_layers=1,
         n_heads=2,
-        d_model=64,
-        d_ffn=64,
+        d_model=32,
+        d_ffn=32,
         moving_avg_window_size=3,
         dropout=0,
         version="Wavelets",  # TODO: Fourier version does not work

tests/imputation/gpvae.py

@@ -45,7 +45,7 @@ class TestGPVAE(unittest.TestCase):
     gp_vae = GPVAE(
         DATA["n_steps"],
         DATA["n_features"],
-        256,
+        32,
         epochs=EPOCHS,
         saving_path=saving_path,
         optimizer=optimizer,

tests/imputation/informer.py

@@ -47,8 +47,8 @@ class TestInformer(unittest.TestCase):
         DATA["n_features"],
         n_layers=2,
         n_heads=2,
-        d_model=128,
-        d_ffn=256,
+        d_model=32,
+        d_ffn=32,
         factor=3,
         dropout=0,
         epochs=EPOCHS,

tests/imputation/mrnn.py

@@ -45,7 +45,7 @@ class TestMRNN(unittest.TestCase):
     mrnn = MRNN(
         DATA["n_steps"],
         DATA["n_features"],
-        256,
+        32,
         epochs=EPOCHS,
         saving_path=saving_path,
         optimizer=optimizer,

tests/imputation/patchtst.py

@@ -46,11 +46,11 @@ class TestPatchTST(unittest.TestCase):
         DATA["n_steps"],
         DATA["n_features"],
         n_layers=2,
-        d_model=256,
-        d_ffn=128,
-        n_heads=4,
-        d_k=64,
-        d_v=64,
+        d_model=64,
+        d_ffn=32,
+        n_heads=2,
+        d_k=16,
+        d_v=16,
         patch_len=16,
         stride=8,
         dropout=0.1,

tests/imputation/saits.py

@@ -46,11 +46,11 @@ class TestSAITS(unittest.TestCase):
         DATA["n_steps"],
         DATA["n_features"],
         n_layers=2,
-        d_model=256,
-        d_ffn=128,
-        n_heads=4,
-        d_k=64,
-        d_v=64,
+        d_model=32,
+        d_ffn=32,
+        n_heads=2,
+        d_k=16,
+        d_v=16,
         dropout=0.1,
         epochs=EPOCHS,
         saving_path=saving_path,

tests/imputation/timesnet.py

@@ -47,8 +47,8 @@ class TestTimesNet(unittest.TestCase):
         DATA["n_features"],
         n_layers=2,
         top_k=3,
-        d_model=128,
-        d_ffn=256,
+        d_model=32,
+        d_ffn=32,
         n_kernels=3,
         dropout=0.1,
         epochs=EPOCHS,

tests/imputation/transformer.py

@@ -46,11 +46,11 @@ class TestTransformer(unittest.TestCase):
         DATA["n_steps"],
         DATA["n_features"],
         n_layers=2,
-        d_model=256,
-        d_ffn=128,
-        n_heads=4,
-        d_k=64,
-        d_v=64,
+        d_model=32,
+        d_ffn=32,
+        n_heads=2,
+        d_k=16,
+        d_v=16,
         dropout=0.1,
         epochs=EPOCHS,
         saving_path=saving_path,

tests/imputation/usgan.py

@@ -46,7 +46,7 @@ class TestUSGAN(unittest.TestCase):
     usgan = USGAN(
         DATA["n_steps"],
         DATA["n_features"],
-        256,
+        32,
         epochs=EPOCHS,
         saving_path=saving_path,
         G_optimizer=G_optimizer,