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/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) 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 61e959f4..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 @@ -546,18 +545,24 @@ 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) self.print_epoch_meters('test') + 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 save_results: - self.logger.info(f'Test results saved to {os.path.join(self.ckpt_save_dir, "test_results.npz")}.') + 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")}.') @@ -591,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') @@ -649,7 +654,6 @@ def test(self, train_epoch: Optional[int] = None, save_metrics: bool = False, sa 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. @@ -919,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 27ed776b..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 @@ -99,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._prediction_memmap = None + self._target_memmap = None + def build_scaler(self, cfg: Dict): """Build scaler. @@ -178,6 +183,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}') + # 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}') def build_train_dataset(self, cfg: Dict): """Build the training dataset. @@ -360,12 +369,14 @@ 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) - self.update_epoch_meter('train/loss', loss.item()) + 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) - 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]): @@ -378,41 +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) - self.update_epoch_meter('val/loss', loss.item()) + 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()) - - 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(), weight) @torch.no_grad() @master_only @@ -425,43 +407,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()) + weight = self._get_metric_weight(forward_return['target']) + self.update_epoch_meter('test/loss', loss.item(), weight) 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() - - 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) + 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, :, :] + 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_h) - 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(), weight) if save_metrics: + metrics_results = {} + 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_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: 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 +505,63 @@ 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 _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: + + self._inputs_memmap = np.memmap(inputs_path, dtype=batch_data['inputs'].dtype, mode='w+', + shape=(total_samples, *batch_data['inputs'].shape[1:])) + self._prediction_memmap = np.memmap(pred_path, dtype=batch_data['prediction'].dtype, mode='w+', + shape=(total_samples, *batch_data['prediction'].shape[1:])) + 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._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._prediction_memmap.flush() + self._target_memmap.flush() + + def _get_metric_weight(self, x: torch.Tensor) -> int: + """ + Get the weight for calculating metrics. + 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. + """ + + 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 valid_num 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'] 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()