feat: 🎸 add the CFG.EVAL.SAVE_RESULTS option (issue #222)

Author: zezhishao

basicts/__init__.py

@@ -1,6 +1,6 @@
 from .launcher import launch_evaluation, launch_training
 from .runners import BaseEpochRunner
 
-__version__ = '0.5.0'
+__version__ = '0.5.1'
 
 __all__ = ['__version__', 'launch_training', 'launch_evaluation', 'BaseEpochRunner']

basicts/runners/base_epoch_runner.py

@@ -87,6 +87,8 @@ def __init__(self, cfg: Dict) -> None:
         self.val_interval = cfg.get('VAL', {}).get('INTERVAL', 1)
         self.test_interval = cfg.get('TEST', {}).get('INTERVAL', 1)
 
+        self.save_results = cfg.get('EVAL', {}).get('SAVE_RESULTS', False)
+
         # create checkpoint save dir
         if not os.path.isdir(self.ckpt_save_dir):
             os.makedirs(self.ckpt_save_dir)
@@ -691,7 +693,7 @@ def on_training_end(self, cfg: Dict, train_epoch: Optional[int] = None):
             )
             self.logger.info('Evaluating the best model on the test set.')
             self.load_model(ckpt_path=best_model_path, strict=True)
-            self.test_pipeline(cfg=cfg, train_epoch=train_epoch, save_metrics=True, save_results=True)
+            self.test_pipeline(cfg=cfg, train_epoch=train_epoch, save_metrics=True, save_results=self.save_results)
 
     # endregion Hook Functions
 

basicts/runners/base_iteration_runner.py

@@ -86,6 +86,8 @@ def __init__(self, cfg: Dict) -> None:
         self.val_data_loader = None
         self.test_data_loader = None
 
+        self.save_results = cfg.get('EVAL', {}).get('SAVE_RESULTS', False)
+
         # declare meter pool
         self.meter_pool = None
 
@@ -574,7 +576,7 @@ def on_training_end(self, cfg: Dict, train_iteration: Optional[int] = None) -> N
             )
             self.logger.info('Evaluating the best model on the test set.')
             self.load_model(ckpt_path=best_model_path, strict=True)
-            self.test_pipeline(cfg=cfg, train_iteration=train_iteration, save_metrics=True, save_results=True)
+            self.test_pipeline(cfg=cfg, train_iteration=train_iteration, save_metrics=True, save_results=self.save_results)
 
     def get_ckpt_path(self, iteration: int) -> str:
         """Get checkpoint path.

examples/complete_config.py

@@ -229,3 +229,4 @@
 # This is a common setting in spatiotemporal forecasting. For long-sequence predictions, it is recommended to keep HORIZONS set to the default value [] to avoid confusion.
 CFG.EVAL.HORIZONS = []
 CFG.EVAL.USE_GPU = True # Whether to use GPU for evaluation. Default: True
+CFG.EVAL.SAVE_RESULTS = False # Whether to save evaluation results in a numpy file. Default: False

examples/complete_config_cn.py

@@ -226,3 +226,4 @@
 # 这是时空预测中的常见配置。对于长序列预测，建议将 HORIZONS 保持为默认值 []，以避免引发误解。
 CFG.EVAL.HORIZONS = []
 CFG.EVAL.USE_GPU = True # 是否在评估时使用 GPU。默认值：True
+CFG.EVAL.SAVE_RESULTS = False # 是否将评估结果保存为一个numpy文件。 默认值：False