clarena.utils.metrics

  1"""The submodule in `utils` for plotting utils."""
  2
  3__all__ = ["MeanMetricBatch", "HATNetworkCapacity"]
  4
  5from typing import Any
  6
  7import torch
  8from torch import Tensor
  9from torchmetrics.aggregation import BaseAggregator
 10
 11
 12class MeanMetricBatch(BaseAggregator):
 13    r"""A TorchMetrics metric to calculate the mean of metrics across data batches.
 14
 15    This is used for accumulated metrics in deep learning. See [here](https://pengxiang-wang.com/posts/continual-learning-metrics.html#nte-accumulate) for more details.
 16    """
 17
 18    def __init__(self, nan_strategy: str | float = "error", **kwargs: Any) -> None:
 19        r"""Initialise the metric. Add state variables."""
 20        BaseAggregator.__init__(
 21            self,
 22            "sum",
 23            torch.tensor(0.0, dtype=torch.get_default_dtype()),
 24            nan_strategy,
 25            state_name="sum",
 26            **kwargs,
 27        )
 28
 29        self.sum: Tensor
 30        r"""State variable created by `super().__init__()` to store the sum of the metric values till this batch."""
 31
 32        self.add_state(
 33            "num",
 34            default=torch.tensor(0, dtype=torch.get_default_dtype()),
 35            dist_reduce_fx="sum",
 36        )
 37        self.num: Tensor
 38        r"""State variable created by `add_state()` to store the number of the data till this batch."""
 39
 40    def update(self, value: torch.Tensor, batch_size: int) -> None:
 41        r"""Update and accumulate the sum of metric value and num of the data till this batch from the batch.
 42
 43        **Args:**
 44        - **val** (`torch.Tensor`): the metric value of the batch to update the sum.
 45        - **batch_size** (`int`): the value to update the num, which is the batch size.
 46        """
 47
 48        value = torch.as_tensor(value, dtype=self.dtype, device=self.device)
 49        batch_size = torch.as_tensor(batch_size, dtype=self.dtype, device=self.device)
 50
 51        self.sum += value * batch_size
 52        self.num += batch_size
 53
 54    def compute(self) -> Tensor:
 55        r"""Compute this mean metric value till this batch.
 56
 57        **Returns:**
 58        - **mean** (`Tensor`): the calculated mean result.
 59        """
 60        return self.sum / self.num
 61
 62
 63class HATNetworkCapacity(BaseAggregator):
 64    r"""A torchmetrics metric to calculate the network capacity of HAT (Hard Attention to the Task) algorithm.
 65
 66    Network capacity is defined as the average adjustment rate over all paramaters. See chapter 4.1 in [AdaHAT paper](https://link.springer.com/chapter/10.1007/978-3-031-70352-2_9).
 67    """
 68
 69    def __init__(self, nan_strategy: str | float = "error", **kwargs: Any) -> None:
 70        r"""Initialise the HAT network capacity metric. Add state variables."""
 71        BaseAggregator.__init__(
 72            self,
 73            "sum",
 74            torch.tensor(0.0, dtype=torch.get_default_dtype()),
 75            nan_strategy,
 76            state_name="sum_adjustment_rate",
 77            **kwargs,
 78        )
 79        self.sum_adjustment_rate: Tensor
 80        r"""State variable created by `add_state()` to store the sum of the adjustment rate values till this layer."""
 81
 82        self.add_state("num_params", default=torch.tensor(0), dist_reduce_fx="sum")
 83        self.num_params: Tensor
 84        r"""State variable created by `add_state()` to store the number of the parameters till this layer."""
 85
 86    def update(
 87        self, adjustment_rate_weight: Tensor, adjustment_rate_bias: Tensor
 88    ) -> None:
 89        r"""Update and accumulate the sum of adjustment rate values till this layer from the layer.
 90
 91        **Args:**
 92        - **adjustment_rate_weight** (`Tensor`): the adjustment rate values of the weight matrix of the layer.
 93        - **adjustment_rate_bias** (`Tensor`): the adjustment rate values of the bias vector of the layer.
 94        """
 95        adjustment_rate_weight = torch.as_tensor(
 96            adjustment_rate_weight, dtype=self.dtype, device=self.device
 97        )
 98        adjustment_rate_bias = torch.as_tensor(
 99            adjustment_rate_bias, dtype=self.dtype, device=self.device
100        )
101
102        self.sum_adjustment_rate += (
103            adjustment_rate_weight.sum() + adjustment_rate_bias.sum()
104        )
105        self.num_params += adjustment_rate_weight.numel() + adjustment_rate_bias.numel()
106
107    def compute(self) -> Tensor:
108        r"""Compute this HAT network capacity till this layer.
109
110        **Returns:**
111        - **network_capacity** (`Tensor`): the calculated network capacity result.
112        """
113
114        return self.sum_adjustment_rate / self.num_params