fix: use the time interval to calculate the TESTAM vocab szie parameter

LibCity · Jun 24, 2024 · 766c164 · 766c164
1 parent 021225b
commit 766c164
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/libcity/config/model/traffic_state_pred/TESTAM.json b/libcity/config/model/traffic_state_pred/TESTAM.json
@@ -24,6 +24,5 @@
   "hidden_size": 32,
   "layers": 3,
   "is_quantile": true,
-  "quantile": 0.7,
-  "vocab_size": 288
+  "quantile": 0.7
 }
diff --git a/libcity/model/traffic_speed_prediction/TESTAM.py b/libcity/model/traffic_speed_prediction/TESTAM.py
@@ -566,7 +566,7 @@ def __init__(self, config, data_feature):
         layers = config.get("layers", 3)
         self.is_quantile = config.get("is_quantile", False)
         self.quantile = config.get("quantile", 0.7)
-        self.vocab_size = config.get("vocab_size", 288)
+        self.vocab_size = 24 * 60 * 60 // config.get("time_intervals", 300)
 
         self.dropout = dropout
         self.prob_mul = prob_mul