improve_speed_tests

Nixtla · Oct 8, 2024 · ccf8b2d · ccf8b2d
1 parent 6c3b2af
commit ccf8b2d
Showing 1 changed file with 21 additions and 31 deletions.
diff --git a/action_files/test_models/src/models2.py b/action_files/test_models/src/models2.py
@@ -2,35 +2,39 @@
 import time
 
 import fire
-import numpy as np
+# import numpy as np
 import pandas as pd
-import pytorch_lightning as pl
-import torch
+# import pytorch_lightning as pl
+# import torch
 
-import neuralforecast
+# import neuralforecast
 from neuralforecast.core import NeuralForecast
 
 from neuralforecast.models.gru import GRU
-from neuralforecast.models.rnn import RNN
-from neuralforecast.models.tcn import TCN
+# from neuralforecast.models.rnn import RNN
+# from neuralforecast.models.tcn import TCN
 from neuralforecast.models.lstm import LSTM
 from neuralforecast.models.dilated_rnn import DilatedRNN
-from neuralforecast.models.deepar import DeepAR
-from neuralforecast.models.mlp import MLP
-from neuralforecast.models.nhits import NHITS
-from neuralforecast.models.nbeats import NBEATS
+# from neuralforecast.models.deepar import DeepAR
+# from neuralforecast.models.mlp import MLP
+# from neuralforecast.models.nhits import NHITS
+# from neuralforecast.models.nbeats import NBEATS
 from neuralforecast.models.nbeatsx import NBEATSx
-from neuralforecast.models.tft import TFT
-from neuralforecast.models.vanillatransformer import VanillaTransformer
-from neuralforecast.models.informer import Informer
-from neuralforecast.models.autoformer import Autoformer
+# from neuralforecast.models.tft import TFT
+# from neuralforecast.models.vanillatransformer import VanillaTransformer
+# from neuralforecast.models.informer import Informer
+# from neuralforecast.models.autoformer import Autoformer
 from neuralforecast.models.patchtst import PatchTST
 
 from neuralforecast.auto import (
-    AutoMLP, AutoNHITS, AutoNBEATS, AutoDilatedRNN, AutoTFT
+    # AutoMLP, 
+    AutoNHITS, 
+    AutoNBEATS, 
+    # AutoDilatedRNN, 
+    # AutoTFT
 )
 
-from neuralforecast.losses.pytorch import SMAPE, MAE
+from neuralforecast.losses.pytorch import MAE
 from ray import tune
 
 from src.data import get_data
@@ -49,23 +53,9 @@ def main(dataset: str = 'M3', group: str = 'Monthly') -> None:
         "scaler_type": "minmax1",
         "random_seed": tune.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
     }
-    config = {
-        "hidden_size": tune.choice([256, 512]),
-        "num_layers": tune.choice([2, 4]),
-        "input_size": tune.choice([2 * horizon]),
-        "max_steps": 1000,
-        "val_check_steps": 300,
-        "scaler_type": "minmax1",
-        "random_seed": tune.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
-    }
-    config_drnn = {'input_size': tune.choice([2 * horizon]),
-                   'encoder_hidden_size': tune.choice([124]),
-                   "max_steps": 300,
-                   "val_check_steps": 100,
-                   "random_seed": tune.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),}
     models = [
         LSTM(h=horizon, input_size=2 * horizon, encoder_hidden_size=50, max_steps=300),
-        DilatedRNN(h=horizon, input_size=2 * horizon, encoder_hidden_size=50, max_steps=300),
+        DilatedRNN(h=horizon, input_size=2 * horizon, encoder_hidden_size=16, max_steps=300),
         GRU(h=horizon, input_size=2 * horizon, encoder_hidden_size=50, max_steps=300),
         AutoNBEATS(h=horizon, loss=MAE(), config=config_nbeats, num_samples=2, cpus=1),
         AutoNHITS(h=horizon, loss=MAE(), config=config_nbeats, num_samples=2, cpus=1),