From 6f832b215818b2e40d21c4c99e018ca50422a387 Mon Sep 17 00:00:00 2001
From: Zezhi Shao <864453277@qq.com>
Date: Wed, 23 Oct 2024 22:31:16 +0800
Subject: [PATCH 01/14] Move the data pre- and post-processing methods to
 simple_tsf_runner.

---
 basicts/runners/base_tsf_runner.py            | 46 ---------
 .../runners/runner_zoo/simple_tsf_runner.py   | 98 ++++++++++++++-----
 2 files changed, 71 insertions(+), 73 deletions(-)

diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index 62afe412..ba40cc37 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -325,46 +325,6 @@ def metric_forward(self, metric_func, args: Dict) -> torch.Tensor:
             raise TypeError(f'Unknown metric type: {type(metric_func)}')
         return metric_item
 
-    def preprocessing(self, input_data: Dict) -> Dict:
-        """Preprocess data.
-
-        Args:
-            input_data (Dict): Dictionary containing data to be processed.
-
-        Returns:
-            Dict: Processed data.
-        """
-
-        if self.scaler is not None:
-            input_data['target'] = self.scaler.transform(input_data['target'])
-            input_data['inputs'] = self.scaler.transform(input_data['inputs'])
-        # TODO: add more preprocessing steps as needed.
-        return input_data
-
-    def postprocessing(self, input_data: Dict) -> Dict:
-        """Postprocess data.
-
-        Args:
-            input_data (Dict): Dictionary containing data to be processed.
-
-        Returns:
-            Dict: Processed data.
-        """
-
-        # rescale data
-        if self.scaler is not None and self.scaler.rescale:
-            input_data['prediction'] = self.scaler.inverse_transform(input_data['prediction'])
-            input_data['target'] = self.scaler.inverse_transform(input_data['target'])
-            input_data['inputs'] = self.scaler.inverse_transform(input_data['inputs'])
-
-        # subset forecasting
-        if self.target_time_series is not None:
-            input_data['target'] = input_data['target'][:, :, self.target_time_series, :]
-            input_data['prediction'] = input_data['prediction'][:, :, self.target_time_series, :]
-
-        # TODO: add more postprocessing steps as needed.
-        return input_data
-
     def train_iters(self, epoch: int, iter_index: int, data: Union[torch.Tensor, Tuple]) -> torch.Tensor:
         """Training iteration process.
 
@@ -378,9 +338,7 @@ def train_iters(self, epoch: int, iter_index: int, data: Union[torch.Tensor, Tup
         """
 
         iter_num = (epoch - 1) * self.iter_per_epoch + iter_index
-        data = self.preprocessing(data)
         forward_return = self.forward(data=data, epoch=epoch, iter_num=iter_num, train=True)
-        forward_return = self.postprocessing(forward_return)
 
         if self.cl_param:
             cl_length = self.curriculum_learning(epoch=epoch)
@@ -402,9 +360,7 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]):
             data (Union[torch.Tensor, Tuple]): Data provided by DataLoader.
         """
 
-        data = self.preprocessing(data)
         forward_return = self.forward(data=data, epoch=None, iter_num=iter_index, train=False)
-        forward_return = self.postprocessing(forward_return)
         loss = self.metric_forward(self.loss, forward_return)
         self.update_epoch_meter('val/loss', loss.item())
 
@@ -456,9 +412,7 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
         prediction, target, inputs = [], [], []
 
         for data in tqdm(self.test_data_loader):
-            data = self.preprocessing(data)
             forward_return = self.forward(data, epoch=None, iter_num=None, train=False)
-            forward_return = self.postprocessing(forward_return)
 
             loss = self.metric_forward(self.loss, forward_return)
             self.update_epoch_meter('test/loss', loss.item())
diff --git a/basicts/runners/runner_zoo/simple_tsf_runner.py b/basicts/runners/runner_zoo/simple_tsf_runner.py
index 4b504ca5..8fb9974c 100644
--- a/basicts/runners/runner_zoo/simple_tsf_runner.py
+++ b/basicts/runners/runner_zoo/simple_tsf_runner.py
@@ -18,48 +18,45 @@ def __init__(self, cfg: Dict):
         self.forward_features = cfg['MODEL'].get('FORWARD_FEATURES', None)
         self.target_features = cfg['MODEL'].get('TARGET_FEATURES', None)
 
-    def select_input_features(self, data: torch.Tensor) -> torch.Tensor:
-        """
-        Selects input features based on the forward features specified in the configuration.
+    def preprocessing(self, input_data: Dict) -> Dict:
+        """Preprocess data.
 
         Args:
-            data (torch.Tensor): Input history data with shape [B, L, N, C1].
+            input_data (Dict): Dictionary containing data to be processed.
 
         Returns:
-            torch.Tensor: Data with selected features with shape [B, L, N, C2].
+            Dict: Processed data.
         """
 
-        if self.forward_features is not None:
-            data = data[:, :, :, self.forward_features]
-        return data
+        if self.scaler is not None:
+            input_data['target'] = self.scaler.transform(input_data['target'])
+            input_data['inputs'] = self.scaler.transform(input_data['inputs'])
+        # TODO: add more preprocessing steps as needed.
+        return input_data
 
-    def select_target_features(self, data: torch.Tensor) -> torch.Tensor:
-        """
-        Selects target features based on the target features specified in the configuration.
+    def postprocessing(self, input_data: Dict) -> Dict:
+        """Postprocess data.
 
         Args:
-            data (torch.Tensor): Model prediction data with shape [B, L, N, C1].
+            input_data (Dict): Dictionary containing data to be processed.
 
         Returns:
-            torch.Tensor: Data with selected target features and shape [B, L, N, C2].
-        """
-
-        data = data[:, :, :, self.target_features]
-        return data
-
-    def select_target_time_series(self, data: torch.Tensor) -> torch.Tensor:
+            Dict: Processed data.
         """
-        Select target time series based on the target time series specified in the configuration.
 
-        Args:
-            data (torch.Tensor): Model prediction data with shape [B, L, N1, C].
+        # rescale data
+        if self.scaler is not None and self.scaler.rescale:
+            input_data['prediction'] = self.scaler.inverse_transform(input_data['prediction'])
+            input_data['target'] = self.scaler.inverse_transform(input_data['target'])
+            input_data['inputs'] = self.scaler.inverse_transform(input_data['inputs'])
 
-        Returns:
-            torch.Tensor: Data with selected target time series and shape [B, L, N2, C].
-        """
+        # subset forecasting
+        if self.target_time_series is not None:
+            input_data['target'] = input_data['target'][:, :, self.target_time_series, :]
+            input_data['prediction'] = input_data['prediction'][:, :, self.target_time_series, :]
 
-        data = data[:, :, self.target_time_series, :]
-        return data
+        # TODO: add more postprocessing steps as needed.
+        return input_data
 
     def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bool = True, **kwargs) -> Dict:
         """
@@ -81,6 +78,8 @@ def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bo
             AssertionError: If the shape of the model output does not match [B, L, N].
         """
 
+        data = self.preprocessing(data)
+
         # Preprocess input data
         future_data, history_data = data['target'], data['inputs']
         history_data = self.to_running_device(history_data)  # Shape: [B, L, N, C]
@@ -111,4 +110,49 @@ def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bo
         assert list(model_return['prediction'].shape)[:3] == [batch_size, length, num_nodes], \
             "The shape of the output is incorrect. Ensure it matches [B, L, N, C]."
 
+        model_return = self.postprocessing(model_return)
+
         return model_return
+
+    def select_input_features(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Selects input features based on the forward features specified in the configuration.
+
+        Args:
+            data (torch.Tensor): Input history data with shape [B, L, N, C1].
+
+        Returns:
+            torch.Tensor: Data with selected features with shape [B, L, N, C2].
+        """
+
+        if self.forward_features is not None:
+            data = data[:, :, :, self.forward_features]
+        return data
+
+    def select_target_features(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Selects target features based on the target features specified in the configuration.
+
+        Args:
+            data (torch.Tensor): Model prediction data with shape [B, L, N, C1].
+
+        Returns:
+            torch.Tensor: Data with selected target features and shape [B, L, N, C2].
+        """
+
+        data = data[:, :, :, self.target_features]
+        return data
+
+    def select_target_time_series(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Select target time series based on the target time series specified in the configuration.
+
+        Args:
+            data (torch.Tensor): Model prediction data with shape [B, L, N1, C].
+
+        Returns:
+            torch.Tensor: Data with selected target time series and shape [B, L, N2, C].
+        """
+
+        data = data[:, :, self.target_time_series, :]
+        return data

From 9adf8c03fcabcf286f956e9df4232d30af2e2afc Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 21 Mar 2025 21:14:56 +0800
Subject: [PATCH 02/14] Remove duplicate code

---
 .../runners/runner_zoo/simple_tsf_runner.py   | 96 -------------------
 1 file changed, 96 deletions(-)

diff --git a/basicts/runners/runner_zoo/simple_tsf_runner.py b/basicts/runners/runner_zoo/simple_tsf_runner.py
index 3f65c2b1..d8372673 100644
--- a/basicts/runners/runner_zoo/simple_tsf_runner.py
+++ b/basicts/runners/runner_zoo/simple_tsf_runner.py
@@ -61,102 +61,6 @@ def postprocessing(self, input_data: Dict) -> Dict:
         # TODO: add more postprocessing steps as needed.
         return input_data
 
-    def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bool = True, **kwargs) -> Dict:
-        """
-        Performs the forward pass for training, validation, and testing. 
-
-        Args:
-            data (Dict): A dictionary containing 'target' (future data) and 'inputs' (history data) (normalized by self.scaler).
-            epoch (int, optional): Current epoch number. Defaults to None.
-            iter_num (int, optional): Current iteration number. Defaults to None.
-            train (bool, optional): Indicates whether the forward pass is for training. Defaults to True.
-
-        Returns:
-            Dict: A dictionary containing the keys:
-                  - 'inputs': Selected input features.
-                  - 'prediction': Model predictions.
-                  - 'target': Selected target features.
-
-        Raises:
-            AssertionError: If the shape of the model output does not match [B, L, N].
-        """
-
-        data = self.preprocessing(data)
-
-        # Preprocess input data
-        future_data, history_data = data['target'], data['inputs']
-        history_data = self.to_running_device(history_data)  # Shape: [B, L, N, C]
-        future_data = self.to_running_device(future_data)    # Shape: [B, L, N, C]
-        batch_size, length, num_nodes, _ = future_data.shape
-
-        # Select input features
-        history_data = self.select_input_features(history_data)
-        future_data_4_dec = self.select_input_features(future_data)
-
-        if not train:
-            # For non-training phases, use only temporal features
-            future_data_4_dec[..., 0] = torch.empty_like(future_data_4_dec[..., 0])
-
-        # Forward pass through the model
-        model_return = self.model(history_data=history_data, future_data=future_data_4_dec,
-                                  batch_seen=iter_num, epoch=epoch, train=train)
-
-        # Parse model return
-        if isinstance(model_return, torch.Tensor):
-            model_return = {'prediction': model_return}
-        if 'inputs' not in model_return:
-            model_return['inputs'] = self.select_target_features(history_data)
-        if 'target' not in model_return:
-            model_return['target'] = self.select_target_features(future_data)
-
-        # Ensure the output shape is correct
-        assert list(model_return['prediction'].shape)[:3] == [batch_size, length, num_nodes], \
-            "The shape of the output is incorrect. Ensure it matches [B, L, N, C]."
-
-        model_return = self.postprocessing(model_return)
-
-        return model_return
-
-    def preprocessing(self, input_data: Dict) -> Dict:
-        """Preprocess data.
-
-        Args:
-            input_data (Dict): Dictionary containing data to be processed.
-
-        Returns:
-            Dict: Processed data.
-        """
-
-        if self.scaler is not None:
-            input_data['target'] = self.scaler.transform(input_data['target'])
-            input_data['inputs'] = self.scaler.transform(input_data['inputs'])
-        # TODO: add more preprocessing steps as needed.
-        return input_data
-
-    def postprocessing(self, input_data: Dict) -> Dict:
-        """Postprocess data.
-
-        Args:
-            input_data (Dict): Dictionary containing data to be processed.
-
-        Returns:
-            Dict: Processed data.
-        """
-
-        # rescale data
-        if self.scaler is not None and self.scaler.rescale:
-            input_data['prediction'] = self.scaler.inverse_transform(input_data['prediction'])
-            input_data['target'] = self.scaler.inverse_transform(input_data['target'])
-            input_data['inputs'] = self.scaler.inverse_transform(input_data['inputs'])
-
-        # subset forecasting
-        if self.target_time_series is not None:
-            input_data['target'] = input_data['target'][:, :, self.target_time_series, :]
-            input_data['prediction'] = input_data['prediction'][:, :, self.target_time_series, :]
-
-        # TODO: add more postprocessing steps as needed.
-        return input_data
-
     def forward(self, data: Dict, epoch: int = None, iter_num: int = None, train: bool = True, **kwargs) -> Dict:
         """
         Performs the forward pass for training, validation, and testing. 

From 0c4a7741069a3127d92da3eb288855b991419134 Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 10:48:00 +0800
Subject: [PATCH 03/14] Update the testing pipeline to fix potential memory
 overflow issues when testing on large-scale datasets.

---
 basicts/runners/base_epoch_runner.py |  14 +--
 basicts/runners/base_tsf_runner.py   | 165 +++++++++++++++++----------
 2 files changed, 111 insertions(+), 68 deletions(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index 61e959f4..88cd6f67 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -552,16 +552,11 @@ def test_pipeline(self, cfg: Optional[Dict] = None, train_epoch: Optional[int] =
         self.update_epoch_meter('test/time', test_end_time - test_start_time)
 
         self.print_epoch_meters('test')
+
         if train_epoch is not None:
             self.plt_epoch_meters('test', train_epoch // self.test_interval)
 
-        # logging here for intuitiveness
-        if save_results:
-            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results.npz")}.')
-        if save_metrics:
-            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
-
-        self.on_test_end()
+        self.on_test_end(save_metrics)
 
     @torch.no_grad()
     @master_only
@@ -642,14 +637,13 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]) -> None:
 
         raise NotImplementedError()
 
-    def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, save_results: bool = False) -> None:
+    def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False) -> None:
         """
         Define the details of the testing process.
 
         Args:
             train_epoch (int, optional): Current epoch during training. Defaults to None.
             save_metrics (bool, optional): Save the test metrics. Defaults to False.
-            save_results (bool, optional): Save the test results. Defaults to False.
 
         Raises:
             NotImplementedError: Must be implemented in a subclass.
@@ -757,7 +751,7 @@ def on_inference_start(self) -> None:
         pass
 
     @master_only
-    def on_test_end(self) -> None:
+    def on_test_end(self, save_metrics: bool = False) -> None:
         """Callback at the end of testing."""
 
         pass
diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index 27ed776b..5066ba1e 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -9,6 +9,7 @@
 import torch
 from easydict import EasyDict
 from easytorch.utils import master_only
+from easytorch.core.meter_pool import MeterPool
 from tqdm import tqdm
 
 from basicts.data.simple_inference_dataset import TimeSeriesInferenceDataset
@@ -178,6 +179,10 @@ def init_test(self, cfg: Dict):
         self.register_epoch_meter('test/loss', 'test', '{:.4f}')
         for key in self.metrics:
             self.register_epoch_meter(f'test/{key}', 'test', '{:.4f}')
+        
+        for i in self.evaluation_horizons:
+            for key in self.metrics:
+                self.register_epoch_meter(f'test/{key}@h{i+1}', f'test @ horizon {i+1}', '{:.4f}')
 
     def build_train_dataset(self, cfg: Dict):
         """Build the training dataset.
@@ -361,7 +366,8 @@ def train_iters(self, epoch: int, iter_index: int, data: Union[torch.Tensor, Tup
             forward_return['prediction'] = forward_return['prediction'][:, :cl_length, :, :]
             forward_return['target'] = forward_return['target'][:, :cl_length, :, :]
         loss = self.metric_forward(self.loss, forward_return)
-        self.update_epoch_meter('train/loss', loss.item())
+        batch_size = forward_return['target'].shape[0]
+        self.update_epoch_meter('train/loss', loss.item(), batch_size)
 
         for metric_name, metric_func in self.metrics.items():
             metric_item = self.metric_forward(metric_func, forward_return)
@@ -378,41 +384,12 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]):
 
         forward_return = self.forward(data=data, epoch=None, iter_num=iter_index, train=False)
         loss = self.metric_forward(self.loss, forward_return)
-        self.update_epoch_meter('val/loss', loss.item())
+        batch_size = forward_return['target'].shape[0]
+        self.update_epoch_meter('val/loss', loss.item(), batch_size)
 
         for metric_name, metric_func in self.metrics.items():
             metric_item = self.metric_forward(metric_func, forward_return)
-            self.update_epoch_meter(f'val/{metric_name}', metric_item.item())
-
-    def compute_evaluation_metrics(self, returns_all: Dict):
-        """Compute metrics for evaluating model performance during the test process.
-
-        Args:
-            returns_all (Dict): Must contain keys: inputs, prediction, target.
-        """
-
-        metrics_results = {}
-        for i in self.evaluation_horizons:
-            pred = returns_all['prediction'][:, i, :, :]
-            real = returns_all['target'][:, i, :, :]
-
-            metrics_results[f'horizon_{i + 1}'] = {}
-            metric_repr = ''
-            for metric_name, metric_func in self.metrics.items():
-                if metric_name.lower() == 'mase':
-                    continue # MASE needs to be calculated after all horizons
-                metric_item = self.metric_forward(metric_func, {'prediction': pred, 'target': real})
-                metric_repr += f', Test {metric_name}: {metric_item.item():.4f}'
-                metrics_results[f'horizon_{i + 1}'][metric_name] = metric_item.item()
-            self.logger.info(f'Evaluate best model on test data for horizon {i + 1}{metric_repr}')
-
-        metrics_results['overall'] = {}
-        for metric_name, metric_func in self.metrics.items():
-            metric_item = self.metric_forward(metric_func, returns_all)
-            self.update_epoch_meter(f'test/{metric_name}', metric_item.item())
-            metrics_results['overall'][metric_name] = metric_item.item()
-
-        return metrics_results
+            self.update_epoch_meter(f'val/{metric_name}', metric_item.item(), batch_size)
 
     @torch.no_grad()
     @master_only
@@ -425,43 +402,53 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
             save_results (bool): Save the test results. Defaults to False.
         """
 
-        prediction, target, inputs = [], [], []
-
-        for data in tqdm(self.test_data_loader):
+        for batch_idx, data in tqdm(enumerate(self.test_data_loader)):
             forward_return = self.forward(data, epoch=None, iter_num=None, train=False)
 
             loss = self.metric_forward(self.loss, forward_return)
-            self.update_epoch_meter('test/loss', loss.item())
+            batch_size = forward_return['target'].shape[0]
+            self.update_epoch_meter('test/loss', loss.item(), batch_size)
 
             if not self.if_evaluate_on_gpu:
-                forward_return['prediction'] = forward_return['prediction'].detach().cpu()
-                forward_return['target'] = forward_return['target'].detach().cpu()
-                forward_return['inputs'] = forward_return['inputs'].detach().cpu()
+                pred = forward_return['prediction'].detach().cpu()
+                target = forward_return['target'].detach().cpu()
+            else:
+                pred = forward_return['prediction']
+                target = forward_return['target']
+            
+            if save_results:
+                batch_data = {
+                    'prediction': forward_return['prediction'].detach().cpu().numpy(),
+                    'target': forward_return['target'].detach().cpu().numpy(),
+                    'inputs': forward_return['inputs'].detach().cpu().numpy()
+                }
+                self._save_test_results(batch_idx, batch_data)
+
+            # evaluation on specific timesteps
+            for i in self.evaluation_horizons:
+                pred_h = pred[:, i, :, :]
+                target_h = target[:, i, :, :]
+
+                for metric_name, metric_func in self.metrics.items():
+                    if metric_name.lower() == 'mase':
+                        continue  # MASE needs to be calculated after all horizons
+                    metric_val = self.metric_forward(metric_func, {'prediction': pred_h, 'target': target_h})
+                    self.update_epoch_meter(f'test/{metric_name}@h{i+1}', metric_val.item(), batch_size)
 
-            prediction.append(forward_return['prediction'])
-            target.append(forward_return['target'])
-            inputs.append(forward_return['inputs'])
-
-        prediction = torch.cat(prediction, dim=0)
-        target = torch.cat(target, dim=0)
-        inputs = torch.cat(inputs, dim=0)
-
-        returns_all = {'prediction': prediction, 'target': target, 'inputs': inputs}
-        metrics_results = self.compute_evaluation_metrics(returns_all)
-
-        # save
-        if save_results:
-            # save returns_all to self.ckpt_save_dir/test_results.npz
-            test_results = {k: v.cpu().numpy() for k, v in returns_all.items()}
-            np.savez(os.path.join(self.ckpt_save_dir, 'test_results.npz'), **test_results)
+            for metric_name, metric_func in self.metrics.items():
+                metric_item = self.metric_forward(metric_func, {'prediction': pred, 'target': target})
+                self.update_epoch_meter(f'test/{metric_name}', metric_item.item(), batch_size)
 
         if save_metrics:
+            metrics_results = {}
+            metrics_results['overall'] = {k: self.meter_pool.get_avg(f'test/{k}') for k in self.metrics.keys()}
+            for i in self.evaluation_horizons:
+                metrics_results[f'horizon_{i+1}'] = {k: self.meter_pool.get_avg(f'test/{k}@h{i+1}') for k in self.metrics.keys()}
+
             # save metrics_results to self.ckpt_save_dir/test_metrics.json
             with open(os.path.join(self.ckpt_save_dir, 'test_metrics.json'), 'w') as f:
                 json.dump(metrics_results, f, indent=4)
 
-        return returns_all
-
     @torch.no_grad()
     @master_only
     def inference(self, save_result_path: str = '') -> tuple:
@@ -513,3 +500,65 @@ def on_validating_end(self, train_epoch: Optional[int]):
         greater_best = not self.metrics_best == 'min'
         if train_epoch is not None:
             self.save_best_model(train_epoch, 'val/' + self.target_metrics, greater_best=greater_best)
+    
+    @master_only
+    def on_test_end(self, save_metrics: bool = False):
+        """Callback at the end of the test process.
+        """
+
+        if len(self.evaluation_horizons) > 0:
+            self.logger.info(f"Evaluation on horizons: {[h + 1 for h in self.evaluation_horizons]}.")
+            for i in self.evaluation_horizons:
+                self.print_epoch_meters(f'test @ horizon {i+1}')
+
+        # logging here for intuitiveness
+        if self.save_results:
+            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
+        if save_metrics:
+            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
+    
+    @master_only
+    def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray]) -> None:
+
+        """
+        Save the test results to disk.
+        
+        Args:
+            batch_idx (int): The index of the current batch.
+            batch_data (Dict[np.ndarray]): The test results:{
+                'inputs': np.ndarray,
+                'prediction': np.ndarray,
+                'target': np.ndarray,
+            }
+        """
+
+        total_samples = len(self.test_data_loader.dataset)
+        
+        save_dir = os.path.join(self.ckpt_save_dir, "test_results")
+        os.makedirs(save_dir, exist_ok=True)
+        inputs_path = os.path.join(save_dir, "inputs.npy")
+        pred_path = os.path.join(save_dir, "predictions.npy")
+        tgt_path = os.path.join(save_dir, "targets.npy")
+
+        # create memmap files
+        if batch_idx == 0:
+
+            global pred_memmap, tgt_memmap, inputs_memmap
+            inputs_memmap = np.memmap(inputs_path, dtype=batch_data['inputs'].dtype, mode='w+',
+                                    shape=(total_samples, *batch_data['inputs'].shape[1:]))
+            pred_memmap = np.memmap(pred_path, dtype=batch_data['prediction'].dtype, mode='w+',
+                                    shape=(total_samples, *batch_data['prediction'].shape[1:]))
+            tgt_memmap = np.memmap(tgt_path, dtype=batch_data['target'].dtype, mode='w+',
+                                shape=(total_samples, *batch_data['target'].shape[1:]))
+
+        start = batch_idx * batch_data['inputs'].shape[0]
+        end = start + batch_data['inputs'].shape[0]
+
+        inputs_memmap[start:end] = batch_data['inputs']
+        pred_memmap[start:end] = batch_data['prediction']
+        tgt_memmap[start:end] = batch_data['target']
+
+        if batch_idx == (total_samples // batch_data['inputs'].shape[0]):
+            inputs_memmap.flush()
+            pred_memmap.flush()
+            tgt_memmap.flush()

From c96e2c955b5ea046640d9702c770fd2907582c2c Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 11:00:15 +0800
Subject: [PATCH 04/14] Update the testing pipeline to fix potential memory
 overflow issues when testing on large-scale datasets.

---
 basicts/runners/base_epoch_runner.py |  2 +-
 basicts/runners/base_tsf_runner.py   | 35 +++++++++++++++-------------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index 88cd6f67..3b382f68 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -637,7 +637,7 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]) -> None:
 
         raise NotImplementedError()
 
-    def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False) -> None:
+    def test(self, train_epoch: Optional[int] = None, save_results: bool = False, save_metrics: bool = False) -> None:
         """
         Define the details of the testing process.
 
diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index 5066ba1e..84e58e88 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -9,7 +9,6 @@
 import torch
 from easydict import EasyDict
 from easytorch.utils import master_only
-from easytorch.core.meter_pool import MeterPool
 from tqdm import tqdm
 
 from basicts.data.simple_inference_dataset import TimeSeriesInferenceDataset
@@ -100,6 +99,11 @@ def __init__(self, cfg: Dict):
         self.evaluation_horizons = [_ - 1 for _ in cfg.get('EVAL', EasyDict()).get('HORIZONS', [])]
         assert len(self.evaluation_horizons) == 0 or min(self.evaluation_horizons) >= 0, 'The horizon should start counting from 1.'
 
+        # For saving test results
+        self._inputs_memmap = None
+        self._pred_memmap = None
+        self._tgt_memmap = None
+
     def build_scaler(self, cfg: Dict):
         """Build scaler.
 
@@ -507,7 +511,7 @@ def on_test_end(self, save_metrics: bool = False):
         """
 
         if len(self.evaluation_horizons) > 0:
-            self.logger.info(f"Evaluation on horizons: {[h + 1 for h in self.evaluation_horizons]}.")
+            self.logger.info(f'Evaluation on horizons: {[h + 1 for h in self.evaluation_horizons]}.')
             for i in self.evaluation_horizons:
                 self.print_epoch_meters(f'test @ horizon {i+1}')
 
@@ -534,31 +538,30 @@ def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray])
 
         total_samples = len(self.test_data_loader.dataset)
         
-        save_dir = os.path.join(self.ckpt_save_dir, "test_results")
+        save_dir = os.path.join(self.ckpt_save_dir, 'test_results')
         os.makedirs(save_dir, exist_ok=True)
-        inputs_path = os.path.join(save_dir, "inputs.npy")
-        pred_path = os.path.join(save_dir, "predictions.npy")
-        tgt_path = os.path.join(save_dir, "targets.npy")
+        inputs_path = os.path.join(save_dir, 'inputs.npy')
+        pred_path = os.path.join(save_dir, 'predictions.npy')
+        tgt_path = os.path.join(save_dir, 'targets.npy')
 
         # create memmap files
         if batch_idx == 0:
 
-            global pred_memmap, tgt_memmap, inputs_memmap
-            inputs_memmap = np.memmap(inputs_path, dtype=batch_data['inputs'].dtype, mode='w+',
+            self._inputs_memmap = np.memmap(inputs_path, dtype=batch_data['inputs'].dtype, mode='w+',
                                     shape=(total_samples, *batch_data['inputs'].shape[1:]))
-            pred_memmap = np.memmap(pred_path, dtype=batch_data['prediction'].dtype, mode='w+',
+            self._pred_memmap = np.memmap(pred_path, dtype=batch_data['prediction'].dtype, mode='w+',
                                     shape=(total_samples, *batch_data['prediction'].shape[1:]))
-            tgt_memmap = np.memmap(tgt_path, dtype=batch_data['target'].dtype, mode='w+',
+            self._tgt_memmap = np.memmap(tgt_path, dtype=batch_data['target'].dtype, mode='w+',
                                 shape=(total_samples, *batch_data['target'].shape[1:]))
 
         start = batch_idx * batch_data['inputs'].shape[0]
         end = start + batch_data['inputs'].shape[0]
 
-        inputs_memmap[start:end] = batch_data['inputs']
-        pred_memmap[start:end] = batch_data['prediction']
-        tgt_memmap[start:end] = batch_data['target']
+        self._inputs_memmap[start:end] = batch_data['inputs']
+        self._pred_memmap[start:end] = batch_data['prediction']
+        self._tgt_memmap[start:end] = batch_data['target']
 
         if batch_idx == (total_samples // batch_data['inputs'].shape[0]):
-            inputs_memmap.flush()
-            pred_memmap.flush()
-            tgt_memmap.flush()
+            self._inputs_memmap.flush()
+            self._pred_memmap.flush()
+            self._tgt_memmap.flush()

From 7e118592c6010e9fd3ea4870bf7b8947a6fdb7b6 Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 11:11:21 +0800
Subject: [PATCH 05/14] Update the testing pipeline to fix potential memory
 overflow issues when testing on large-scale datasets.

---
 basicts/runners/base_epoch_runner.py | 4 ++--
 basicts/runners/base_tsf_runner.py   | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index 3b382f68..5ececb67 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -546,7 +546,7 @@ def test_pipeline(self, cfg: Optional[Dict] = None, train_epoch: Optional[int] =
         self.model.eval()
 
         # execute the test process
-        self.test(train_epoch=train_epoch, save_results=save_results, save_metrics=save_metrics)
+        self.test(train_epoch=train_epoch, save_metrics=save_metrics, save_results=save_results)
 
         test_end_time = time.time()
         self.update_epoch_meter('test/time', test_end_time - test_start_time)
@@ -637,7 +637,7 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]) -> None:
 
         raise NotImplementedError()
 
-    def test(self, train_epoch: Optional[int] = None, save_results: bool = False, save_metrics: bool = False) -> None:
+    def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, save_results: bool = False,) -> None:
         """
         Define the details of the testing process.
 
diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index 84e58e88..81e0a526 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -183,7 +183,7 @@ def init_test(self, cfg: Dict):
         self.register_epoch_meter('test/loss', 'test', '{:.4f}')
         for key in self.metrics:
             self.register_epoch_meter(f'test/{key}', 'test', '{:.4f}')
-        
+        # Register metrics for each evaluation horizons
         for i in self.evaluation_horizons:
             for key in self.metrics:
                 self.register_epoch_meter(f'test/{key}@h{i+1}', f'test @ horizon {i+1}', '{:.4f}')
@@ -419,7 +419,6 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
             else:
                 pred = forward_return['prediction']
                 target = forward_return['target']
-            
             if save_results:
                 batch_data = {
                     'prediction': forward_return['prediction'].detach().cpu().numpy(),

From 84fdc8d6302530354e5427b5d422bdbc8154310f Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 11:15:55 +0800
Subject: [PATCH 06/14] Update the testing pipeline to fix potential memory
 overflow issues when testing on large-scale datasets.

---
 basicts/runners/base_tsf_runner.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index 81e0a526..e52494dc 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -519,7 +519,7 @@ def on_test_end(self, save_metrics: bool = False):
             self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
         if save_metrics:
             self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
-    
+
     @master_only
     def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray]) -> None:
 
@@ -536,7 +536,7 @@ def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray])
         """
 
         total_samples = len(self.test_data_loader.dataset)
-        
+
         save_dir = os.path.join(self.ckpt_save_dir, 'test_results')
         os.makedirs(save_dir, exist_ok=True)
         inputs_path = os.path.join(save_dir, 'inputs.npy')

From fc165ecd605ac59c1cb0bd9718ed92dae9babfe3 Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 11:18:10 +0800
Subject: [PATCH 07/14] Update the testing pipeline to fix potential memory
 overflow issues when testing on large-scale datasets.

---
 basicts/runners/base_tsf_runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index e52494dc..efb5f329 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -503,7 +503,7 @@ def on_validating_end(self, train_epoch: Optional[int]):
         greater_best = not self.metrics_best == 'min'
         if train_epoch is not None:
             self.save_best_model(train_epoch, 'val/' + self.target_metrics, greater_best=greater_best)
-    
+
     @master_only
     def on_test_end(self, save_metrics: bool = False):
         """Callback at the end of the test process.

From 83c9d69f2b64ca54801a305b6b594b64f44ce622 Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 16:02:02 +0800
Subject: [PATCH 08/14] Fix: update according to review comments.

---
 basicts/runners/base_epoch_runner.py            | 12 +++++++++---
 basicts/runners/runner_zoo/simple_tsf_runner.py |  4 ++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index 5ececb67..c0c73b07 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -556,7 +556,13 @@ def test_pipeline(self, cfg: Optional[Dict] = None, train_epoch: Optional[int] =
         if train_epoch is not None:
             self.plt_epoch_meters('test', train_epoch // self.test_interval)
 
-        self.on_test_end(save_metrics)
+        # logging here for intuitiveness
+        if self.save_results:
+            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
+        if save_metrics:
+            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
+        
+        self.on_test_end()
 
     @torch.no_grad()
     @master_only
@@ -637,7 +643,7 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]) -> None:
 
         raise NotImplementedError()
 
-    def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, save_results: bool = False,) -> None:
+    def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, save_results: bool = False) -> None:
         """
         Define the details of the testing process.
 
@@ -751,7 +757,7 @@ def on_inference_start(self) -> None:
         pass
 
     @master_only
-    def on_test_end(self, save_metrics: bool = False) -> None:
+    def on_test_end(self) -> None:
         """Callback at the end of testing."""
 
         pass
diff --git a/basicts/runners/runner_zoo/simple_tsf_runner.py b/basicts/runners/runner_zoo/simple_tsf_runner.py
index af97861a..5f616b51 100644
--- a/basicts/runners/runner_zoo/simple_tsf_runner.py
+++ b/basicts/runners/runner_zoo/simple_tsf_runner.py
@@ -92,13 +92,13 @@ def forward(self, data: Dict, epoch: Optional[int] = None, iter_num: Optional[in
         # Select input features
         history_data = self.select_input_features(history_data)
         future_data_4_dec = self.select_input_features(future_data)
-        
+
         if not train:
             # For non-training phases, use only temporal features
             future_data_4_dec[..., 0] = torch.empty_like(future_data_4_dec[..., 0])
 
         # Forward pass through the model
-        model_return = self.model(history_data=history_data, future_data=future_data_4_dec, 
+        model_return = self.model(history_data=history_data, future_data=future_data_4_dec,
                                   batch_seen=iter_num, epoch=epoch, train=train)
 
         # Parse model return

From bfea179092edd57b8f6de69d4bf58c312e096e63 Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Fri, 15 Aug 2025 16:05:51 +0800
Subject: [PATCH 09/14] Fix: update according to review comments.

---
 basicts/runners/base_epoch_runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index c0c73b07..0eeeffec 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -561,7 +561,7 @@ def test_pipeline(self, cfg: Optional[Dict] = None, train_epoch: Optional[int] =
             self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
         if save_metrics:
             self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
-        
+
         self.on_test_end()
 
     @torch.no_grad()

From 05aa1b7e18c05f15ef1af5a2af84fffd5a368d9b Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Mon, 18 Aug 2025 16:48:33 +0800
Subject: [PATCH 10/14] =?UTF-8?q?Fix:=20Corrected=20the=20bug=20where=20me?=
 =?UTF-8?q?trics=20were=20averaged=20without=20accounting=20for=20missing?=
 =?UTF-8?q?=20values=E2=80=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 basicts/runners/base_epoch_runner.py |  9 +++-
 basicts/runners/base_tsf_runner.py   | 70 +++++++++++++++-------------
 2 files changed, 44 insertions(+), 35 deletions(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index 0eeeffec..0155d98f 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -556,11 +556,16 @@ def test_pipeline(self, cfg: Optional[Dict] = None, train_epoch: Optional[int] =
         if train_epoch is not None:
             self.plt_epoch_meters('test', train_epoch // self.test_interval)
 
+        if len(self.evaluation_horizons) > 0:
+            self.logger.info(f'Evaluation on horizons: {[h + 1 for h in self.evaluation_horizons]}.')
+            for i in self.evaluation_horizons:
+                self.print_epoch_meters(f'test @ horizon {i+1}')
+
         # logging here for intuitiveness
         if self.save_results:
-            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
+            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, 'test_results')}.')
         if save_metrics:
-            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
+            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, 'test_metrics.json')}.')
 
         self.on_test_end()
 
diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index efb5f329..bc04bf14 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -101,8 +101,8 @@ def __init__(self, cfg: Dict):
 
         # For saving test results
         self._inputs_memmap = None
-        self._pred_memmap = None
-        self._tgt_memmap = None
+        self._prediction_memmap = None
+        self._target_memmap = None
 
     def build_scaler(self, cfg: Dict):
         """Build scaler.
@@ -369,9 +369,10 @@ def train_iters(self, epoch: int, iter_index: int, data: Union[torch.Tensor, Tup
             cl_length = self.curriculum_learning(epoch=epoch)
             forward_return['prediction'] = forward_return['prediction'][:, :cl_length, :, :]
             forward_return['target'] = forward_return['target'][:, :cl_length, :, :]
+
         loss = self.metric_forward(self.loss, forward_return)
-        batch_size = forward_return['target'].shape[0]
-        self.update_epoch_meter('train/loss', loss.item(), batch_size)
+        weight = self._get_metric_weight(forward_return['target'])
+        self.update_epoch_meter('train/loss', loss.item(), weight)
 
         for metric_name, metric_func in self.metrics.items():
             metric_item = self.metric_forward(metric_func, forward_return)
@@ -388,12 +389,12 @@ def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]):
 
         forward_return = self.forward(data=data, epoch=None, iter_num=iter_index, train=False)
         loss = self.metric_forward(self.loss, forward_return)
-        batch_size = forward_return['target'].shape[0]
-        self.update_epoch_meter('val/loss', loss.item(), batch_size)
+        weight = self._get_metric_weight(forward_return['target'])
+        self.update_epoch_meter('val/loss', loss.item(), weight)
 
         for metric_name, metric_func in self.metrics.items():
             metric_item = self.metric_forward(metric_func, forward_return)
-            self.update_epoch_meter(f'val/{metric_name}', metric_item.item(), batch_size)
+            self.update_epoch_meter(f'val/{metric_name}', metric_item.item(), weight)
 
     @torch.no_grad()
     @master_only
@@ -410,8 +411,8 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
             forward_return = self.forward(data, epoch=None, iter_num=None, train=False)
 
             loss = self.metric_forward(self.loss, forward_return)
-            batch_size = forward_return['target'].shape[0]
-            self.update_epoch_meter('test/loss', loss.item(), batch_size)
+            weight = self._get_metric_weight(forward_return['target'])
+            self.update_epoch_meter('test/loss', loss.item(), weight)
 
             if not self.if_evaluate_on_gpu:
                 pred = forward_return['prediction'].detach().cpu()
@@ -436,11 +437,11 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
                     if metric_name.lower() == 'mase':
                         continue  # MASE needs to be calculated after all horizons
                     metric_val = self.metric_forward(metric_func, {'prediction': pred_h, 'target': target_h})
-                    self.update_epoch_meter(f'test/{metric_name}@h{i+1}', metric_val.item(), batch_size)
+                    self.update_epoch_meter(f'test/{metric_name}@h{i+1}', metric_val.item(), weight)
 
             for metric_name, metric_func in self.metrics.items():
                 metric_item = self.metric_forward(metric_func, {'prediction': pred, 'target': target})
-                self.update_epoch_meter(f'test/{metric_name}', metric_item.item(), batch_size)
+                self.update_epoch_meter(f'test/{metric_name}', metric_item.item(), weight)
 
         if save_metrics:
             metrics_results = {}
@@ -504,22 +505,6 @@ def on_validating_end(self, train_epoch: Optional[int]):
         if train_epoch is not None:
             self.save_best_model(train_epoch, 'val/' + self.target_metrics, greater_best=greater_best)
 
-    @master_only
-    def on_test_end(self, save_metrics: bool = False):
-        """Callback at the end of the test process.
-        """
-
-        if len(self.evaluation_horizons) > 0:
-            self.logger.info(f'Evaluation on horizons: {[h + 1 for h in self.evaluation_horizons]}.')
-            for i in self.evaluation_horizons:
-                self.print_epoch_meters(f'test @ horizon {i+1}')
-
-        # logging here for intuitiveness
-        if self.save_results:
-            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
-        if save_metrics:
-            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
-
     @master_only
     def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray]) -> None:
 
@@ -548,19 +533,38 @@ def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray])
 
             self._inputs_memmap = np.memmap(inputs_path, dtype=batch_data['inputs'].dtype, mode='w+',
                                     shape=(total_samples, *batch_data['inputs'].shape[1:]))
-            self._pred_memmap = np.memmap(pred_path, dtype=batch_data['prediction'].dtype, mode='w+',
+            self._prediction_memmap = np.memmap(pred_path, dtype=batch_data['prediction'].dtype, mode='w+',
                                     shape=(total_samples, *batch_data['prediction'].shape[1:]))
-            self._tgt_memmap = np.memmap(tgt_path, dtype=batch_data['target'].dtype, mode='w+',
+            self._target_memmap = np.memmap(tgt_path, dtype=batch_data['target'].dtype, mode='w+',
                                 shape=(total_samples, *batch_data['target'].shape[1:]))
 
         start = batch_idx * batch_data['inputs'].shape[0]
         end = start + batch_data['inputs'].shape[0]
 
         self._inputs_memmap[start:end] = batch_data['inputs']
-        self._pred_memmap[start:end] = batch_data['prediction']
-        self._tgt_memmap[start:end] = batch_data['target']
+        self._prediction_memmap[start:end] = batch_data['prediction']
+        self._target_memmap[start:end] = batch_data['target']
 
         if batch_idx == (total_samples // batch_data['inputs'].shape[0]):
             self._inputs_memmap.flush()
-            self._pred_memmap.flush()
-            self._tgt_memmap.flush()
+            self._prediction_memmap.flush()
+            self._target_memmap.flush()
+
+    def _get_metric_weight(self, x: torch.Tensor) -> int:
+        """
+        Get the weight for calculating metrics.
+        1. Since the last batch may be smaller (`drop_last=False`), it is necessary to perform a weighted average based on the batch size.
+        2. Since the number of valid values in each batch may vary, a weighted average based on the valid value count is also required.
+           Valid value count is the total count minus the number of missing values.
+        The weight is the product of the batch size and the valid value count.
+        """
+
+        batch_size = x.shape[0]
+
+        if self.null_val == np.nan:
+            valid_num = (~torch.isnan(x)).sum().item()
+        else:
+            eps = 5e-5
+            valid_num = (~torch.isclose(x, torch.tensor(self.null_val).expand_as(x).to(x.device), atol=eps, rtol=0.0)).sum().item()
+
+        return batch_size * valid_num

From f3b725b01382ee05566164ad6fcb4965753679fa Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Mon, 18 Aug 2025 17:03:32 +0800
Subject: [PATCH 11/14] =?UTF-8?q?Fix:=20Corrected=20the=20bug=20where=20me?=
 =?UTF-8?q?trics=20were=20averaged=20without=20accounting=20for=20missing?=
 =?UTF-8?q?=20values=E2=80=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 basicts/runners/base_epoch_runner.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index 0155d98f..d53d15fd 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -563,9 +563,9 @@ def test_pipeline(self, cfg: Optional[Dict] = None, train_epoch: Optional[int] =
 
         # logging here for intuitiveness
         if self.save_results:
-            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, 'test_results')}.')
+            self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results")}.')
         if save_metrics:
-            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, 'test_metrics.json')}.')
+            self.logger.info(f'Test metrics saved to {os.path.join(self.ckpt_save_dir, "test_metrics.json")}.')
 
         self.on_test_end()
 

From 2d729f2401a575a7ceaeeeba828d9adc3afdec35 Mon Sep 17 00:00:00 2001
From: yisongfu <jokerfulyfock@gmail.com>
Date: Tue, 19 Aug 2025 12:42:01 +0800
Subject: [PATCH 12/14] =?UTF-8?q?Fix:=20Corrected=20the=20bug=20where=20me?=
 =?UTF-8?q?trics=20were=20averaged=20without=20accounting=20for=20missing?=
 =?UTF-8?q?=20values=E2=80=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 basicts/metrics/__init__.py          |  21 +++++-
 basicts/metrics/metric_meter.py      |  76 ++++++++++++++++++++
 basicts/runners/base_epoch_runner.py |   7 +-
 basicts/runners/base_tsf_runner.py   |  21 +++---
 basicts/utils/meter_pool.py          | 104 +++++++++++++++++++++++++++
 5 files changed, 211 insertions(+), 18 deletions(-)
 create mode 100644 basicts/metrics/metric_meter.py
 create mode 100644 basicts/utils/meter_pool.py

diff --git a/basicts/metrics/__init__.py b/basicts/metrics/__init__.py
index 66e465b5..44807634 100644
--- a/basicts/metrics/__init__.py
+++ b/basicts/metrics/__init__.py
@@ -1,6 +1,8 @@
+from .cls_metrics import accuracy, f1_score, precision, recall
 from .corr import masked_corr
 from .mae import masked_mae
 from .mape import masked_mape
+from .metric_meter import AvgMeter, RMSEMeter
 from .mse import masked_mse
 from .r_square import masked_r2
 from .rmse import masked_rmse
@@ -15,17 +17,32 @@
             'WAPE': masked_wape,
             'SMAPE': masked_smape,
             'R2': masked_r2,
-            'CORR': masked_corr
+            'CORR': masked_corr,
+            "accuracy": accuracy,
+            "precision": precision,
+            "recall": recall,
+            "f1": f1_score
             }
 
+METRIC_METER = {
+    'RMSE': RMSEMeter,
+    'default': AvgMeter
+}
+
 __all__ = [
     'masked_mae',
     'masked_mse',
     'masked_rmse',
+    'incremental_masked_rmse',
     'masked_mape',
     'masked_wape',
     'masked_smape',
     'masked_r2',
     'masked_corr',
-    'ALL_METRICS'
+    'accuracy',
+    'precision',
+    'recall',
+    'f1_score',
+    'ALL_METRICS',
+    'METRIC_METER'
 ]
\ No newline at end of file
diff --git a/basicts/metrics/metric_meter.py b/basicts/metrics/metric_meter.py
new file mode 100644
index 00000000..8ebc8012
--- /dev/null
+++ b/basicts/metrics/metric_meter.py
@@ -0,0 +1,76 @@
+class AvgMeter:
+    """Average meter.
+    """
+
+    def __init__(self):
+        self._sum: float = 0.
+        self._count: int = 0
+
+    def reset(self):
+        """Reset counter.
+        """
+
+        self._sum = 0.
+        self._count = 0
+
+    def update(self, value: float, n: int = 1):
+        """Update sum and count.
+
+        Args:
+            value (float): value.
+            n (int): number.
+        """
+
+        self._sum += value * n
+        self._count += n
+
+    @property
+    def value(self) -> float:
+        """Get average value.
+
+        Returns:
+            avg (float)
+        """
+
+        return self._sum / self._count if self._count != 0 else 0
+
+
+class RMSEMeter:
+    """
+    RMSE meter.
+    This meter maintains **MSE** and calculate **RMSE** in the post process.
+    """
+
+    def __init__(self):
+        self._mse: float = 0.
+        self._count: int = 0
+
+    def reset(self):
+        """Reset counter.
+        """
+
+        self._mse = 0.
+        self._count = 0
+
+    def update(self, value: float, n: int = 1):
+        """Update sum and count.
+
+        Args:
+            value (float): value.
+            n (int): number.
+        """
+
+        self._mse += value ** 2 * n
+        self._count += n
+
+    @property
+    def value(self) -> float:
+        """Get average value.
+
+        Returns:
+            avg (float)
+        """
+
+        mse = self._mse / self._count if self._count != 0 else 0
+
+        return mse ** 0.5
diff --git a/basicts/runners/base_epoch_runner.py b/basicts/runners/base_epoch_runner.py
index d53d15fd..2cb17e92 100644
--- a/basicts/runners/base_epoch_runner.py
+++ b/basicts/runners/base_epoch_runner.py
@@ -10,7 +10,6 @@
 from easytorch.core.checkpoint import (backup_last_ckpt, clear_ckpt, load_ckpt,
                                        save_ckpt)
 from easytorch.core.data_loader import build_data_loader, build_data_loader_ddp
-from easytorch.core.meter_pool import MeterPool
 from easytorch.device import to_device
 from easytorch.utils import (TimePredictor, get_local_rank, get_logger,
                              is_master, master_only, set_env)
@@ -22,7 +21,7 @@
 from torch.utils.tensorboard import SummaryWriter
 from tqdm import tqdm
 
-from ..utils import get_dataset_name
+from ..utils import MeterPool, get_dataset_name
 from . import optim
 
 
@@ -597,7 +596,7 @@ def inference_pipeline(self, cfg: Optional[Dict] = None, input_data: Union[str,
         result = self.inference(save_result_path=output_data_file_path)
 
         inference_end_time = time.time()
-        self.update_epoch_meter('inference/time', inference_end_time - inference_start_time)
+        self.update_epoch_meter('inference/time', 'inference', inference_end_time - inference_start_time)
 
         self.print_epoch_meters('inference')
 
@@ -924,7 +923,7 @@ def save_best_model(self, epoch: int, metric_name: str, greater_best: bool = Tru
                 `False` means lower value is best, such as `loss`. Defaults to True.
         """
 
-        metric = self.meter_pool.get_avg(metric_name)
+        metric = self.meter_pool.get_value(metric_name)
         best_metric = self.best_metrics.get(metric_name)
         if best_metric is None or (metric > best_metric if greater_best else metric < best_metric):
             self.best_metrics[metric_name] = metric
diff --git a/basicts/runners/base_tsf_runner.py b/basicts/runners/base_tsf_runner.py
index bc04bf14..a46085e7 100644
--- a/basicts/runners/base_tsf_runner.py
+++ b/basicts/runners/base_tsf_runner.py
@@ -74,7 +74,7 @@ def __init__(self, cfg: Dict):
         # define metrics
         self.metrics = cfg.get('METRICS', {}).get('FUNCS', {
                                                             'MAE': masked_mae, 
-                                                            'RMSE': masked_rmse, 
+                                                            'RMSE': masked_rmse,
                                                             'MAPE': masked_mape, 
                                                             'WAPE': masked_wape, 
                                                             'MSE': masked_mse
@@ -376,7 +376,7 @@ def train_iters(self, epoch: int, iter_index: int, data: Union[torch.Tensor, Tup
 
         for metric_name, metric_func in self.metrics.items():
             metric_item = self.metric_forward(metric_func, forward_return)
-            self.update_epoch_meter(f'train/{metric_name}', metric_item.item())
+            self.update_epoch_meter(f'train/{metric_name}', metric_item.item(), weight)
         return loss
 
     def val_iters(self, iter_index: int, data: Union[torch.Tensor, Tuple]):
@@ -432,12 +432,13 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
             for i in self.evaluation_horizons:
                 pred_h = pred[:, i, :, :]
                 target_h = target[:, i, :, :]
+                weight_h = self._get_metric_weight(target_h)
 
                 for metric_name, metric_func in self.metrics.items():
                     if metric_name.lower() == 'mase':
                         continue  # MASE needs to be calculated after all horizons
                     metric_val = self.metric_forward(metric_func, {'prediction': pred_h, 'target': target_h})
-                    self.update_epoch_meter(f'test/{metric_name}@h{i+1}', metric_val.item(), weight)
+                    self.update_epoch_meter(f'test/{metric_name}@h{i+1}', metric_val.item(), weight_h)
 
             for metric_name, metric_func in self.metrics.items():
                 metric_item = self.metric_forward(metric_func, {'prediction': pred, 'target': target})
@@ -445,9 +446,9 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa
 
         if save_metrics:
             metrics_results = {}
-            metrics_results['overall'] = {k: self.meter_pool.get_avg(f'test/{k}') for k in self.metrics.keys()}
+            metrics_results['overall'] = {k: self.meter_pool.get_value(f'test/{k}') for k in self.metrics.keys()}
             for i in self.evaluation_horizons:
-                metrics_results[f'horizon_{i+1}'] = {k: self.meter_pool.get_avg(f'test/{k}@h{i+1}') for k in self.metrics.keys()}
+                metrics_results[f'horizon_{i+1}'] = {k: self.meter_pool.get_value(f'test/{k}@h{i+1}') for k in self.metrics.keys()}
 
             # save metrics_results to self.ckpt_save_dir/test_metrics.json
             with open(os.path.join(self.ckpt_save_dir, 'test_metrics.json'), 'w') as f:
@@ -553,18 +554,14 @@ def _save_test_results(self, batch_idx: int, batch_data: Dict[str, np.ndarray])
     def _get_metric_weight(self, x: torch.Tensor) -> int:
         """
         Get the weight for calculating metrics.
-        1. Since the last batch may be smaller (`drop_last=False`), it is necessary to perform a weighted average based on the batch size.
-        2. Since the number of valid values in each batch may vary, a weighted average based on the valid value count is also required.
-           Valid value count is the total count minus the number of missing values.
-        The weight is the product of the batch size and the valid value count.
+        Since the number of valid values in each batch may vary, it is necessary to perform a weighted average based on the valid value count.
+        The valid value count is the total count minus the number of missing values.
         """
 
-        batch_size = x.shape[0]
-
         if self.null_val == np.nan:
             valid_num = (~torch.isnan(x)).sum().item()
         else:
             eps = 5e-5
             valid_num = (~torch.isclose(x, torch.tensor(self.null_val).expand_as(x).to(x.device), atol=eps, rtol=0.0)).sum().item()
 
-        return batch_size * valid_num
+        return valid_num
diff --git a/basicts/utils/meter_pool.py b/basicts/utils/meter_pool.py
new file mode 100644
index 00000000..761aaa2d
--- /dev/null
+++ b/basicts/utils/meter_pool.py
@@ -0,0 +1,104 @@
+import logging
+from typing import Any, Dict, Tuple, Union
+
+from torch.utils.tensorboard import SummaryWriter
+
+from ..metrics import METRIC_METER
+
+
+class MeterPool:
+    """Meter container
+    """
+
+    def __init__(self):
+        self._pool: Dict[str, Dict[str, Any]] = {}
+
+    def register(self, name: str, meter_type: str, fmt: str = '{:f}', plt: bool = True):
+        """Add a meter to meter pool.
+        Args:
+            name (str): meter name.
+            meter_type (str): meter type.
+            fmt (str): meter output format.
+            plt (bool): set ```True``` to plot it in tensorboard
+                when calling ```plt_meters```.
+        """
+
+        if name in self._pool:
+            raise ValueError(f'Meter {name} already existed.')
+
+        # name: type/metric or type/metric@h{i}
+        metric = name.split('/')[1].split('@')[0] # get the metric name
+        handle_meter = 'default' if metric not in METRIC_METER else metric
+
+        self._pool[name] = {
+            'meter': METRIC_METER[handle_meter](),
+            'index': len(self._pool.keys()),
+            'format': fmt,
+            'type': meter_type,
+            'plt': plt
+        }
+
+    def update(self, name: str, value: Union[float, Tuple[float]] , n: int = 1):
+        """Update average meter.
+
+        Args:
+            name (str): meter name.
+            value (Union[float, Tuple[float]]): value.
+            n: (int): num.
+        """
+
+        self._pool[name]['meter'].update(value, n)
+
+    def get_value(self, name: str) -> float:
+        """Get value.
+
+        Args:
+            name (str): meter name.
+
+        Returns:
+            avg (float)
+        """
+
+        return self._pool[name]['meter'].value
+
+    def print_meters(self, meter_type: str, logger: logging.Logger = None):
+        """Print the specified type of meters.
+
+        Args:
+            meter_type (str): meter type
+            logger (logging.Logger): logger
+        """
+
+        print_list = []
+        for i in range(len(self._pool.keys())):
+            for name, value in self._pool.items():
+                if value['index'] == i and value['type'] == meter_type:
+                    print_list.append(
+                        ('{}: ' + value['format']).format(name, value['meter'].value)
+                    )
+        print_str = 'Result <{}>: [{}]'.format(meter_type, ', '.join(print_list))
+        if logger is None:
+            print(print_str)
+        else:
+            logger.info(print_str)
+
+    def plt_meters(self, meter_type: str, step: int, tensorboard_writer: SummaryWriter):
+        """Plot the specified type of meters in tensorboard.
+
+        Args:
+            meter_type (str): meter type.
+            step (int): Global step value to record
+            tensorboard_writer (SummaryWriter): tensorboard SummaryWriter
+        """
+
+        for name, value in self._pool.items():
+            if value['plt'] and value['type'] == meter_type:
+                tensorboard_writer.add_scalar(name, value['meter'].value, global_step=step)
+        tensorboard_writer.flush()
+
+    def reset(self):
+        """Reset all meters.
+        """
+
+        for _, value in self._pool.items():
+            value['meter'].reset()

From cd16eb27fef282ea196c815a0796b65dcf67057a Mon Sep 17 00:00:00 2001
From: Yisong Fu <139831104+yisongfu@users.noreply.github.com>
Date: Tue, 19 Aug 2025 13:23:31 +0800
Subject: [PATCH 13/14] Update __init__.py

---
 basicts/utils/__init__.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/basicts/utils/__init__.py b/basicts/utils/__init__.py
index 66d5b13c..3746cfae 100644
--- a/basicts/utils/__init__.py
+++ b/basicts/utils/__init__.py
@@ -1,5 +1,6 @@
 from .config import get_dataset_name
 from .dataset import InfiniteGenerator
+from .meter_pool import MeterPool
 from .misc import check_nan_inf, clock
 from .misc import partial_func as partial
 from .misc import remove_nan_inf
@@ -12,4 +13,4 @@
            'remove_nan_inf', 'data_transformation_4_xformer',
            'partial', 'get_regular_settings',
            'load_dataset_data', 'load_dataset_desc',
-           'InfiniteGenerator', 'get_dataset_name']
+           'InfiniteGenerator', 'get_dataset_name', 'MeterPool']

From a582fc4bf81695f9c8be952e2b7b342194259290 Mon Sep 17 00:00:00 2001
From: Yisong Fu <139831104+yisongfu@users.noreply.github.com>
Date: Tue, 19 Aug 2025 13:31:08 +0800
Subject: [PATCH 14/14] Add metrics for time series classification

---
 basicts/metrics/cls_metrics.py | 60 ++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 basicts/metrics/cls_metrics.py

diff --git a/basicts/metrics/cls_metrics.py b/basicts/metrics/cls_metrics.py
new file mode 100644
index 00000000..4846fee1
--- /dev/null
+++ b/basicts/metrics/cls_metrics.py
@@ -0,0 +1,60 @@
+import torch
+
+
+def accuracy(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Calculate the accuracy of predictions.
+
+    Args:
+        pred (torch.Tensor): The predicted values as a tensor.
+        target (torch.Tensor): The ground truth values as a tensor with the same shape as `pred`.
+
+    Returns:
+        torch.Tensor: A scalar tensor representing the accuracy.
+    """
+    return (pred == target).float().mean()
+
+def precision(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Calculate the precision of predictions.
+
+    Args:
+        pred (torch.Tensor): The predicted values as a tensor.
+        target (torch.Tensor): The ground truth values as a tensor with the same shape as `pred`.
+
+    Returns:
+        torch.Tensor: A scalar tensor representing the precision.
+    """
+    true_positives = (pred == target).float().sum()
+    false_positives = (pred != target).float().sum()
+    return true_positives / (true_positives + false_positives)
+
+def recall(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Calculate the recall of predictions.
+
+    Args:
+        pred (torch.Tensor): The predicted values as a tensor.
+        target (torch.Tensor): The ground truth values as a tensor with the same shape as `pred`.
+
+    Returns:
+        torch.Tensor: A scalar tensor representing the recall.
+    """
+    true_positives = (pred == target).float().sum()
+    false_negatives = (pred != target).float().sum()
+    return true_positives / (true_positives + false_negatives)
+
+def f1_score(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    """
+    Calculate the F1 score of predictions.
+
+    Args:
+        pred (torch.Tensor): The predicted values as a tensor.
+        target (torch.Tensor): The ground truth values as a tensor with the same shape as `pred`.
+
+    Returns:
+        torch.Tensor: A scalar tensor representing the F1 score.
+    """
+    precision_item = precision(pred, target)
+    recall_item = recall(pred, target)
+    return 2 * (precision_item * recall_item) / (precision_item + recall_item)