Skip to content

Metrics and losses

Metrics

metrics

Evaluation metrics for images, volumes and compressed representations.

Thin wrappers around torchmetrics that score a single sample and return a Python float. Image metrics take (channels, H, W) images, and data_range defaults to the dtype convention of torchvision and scikit-image: 255 for uint8 images and 1 for floating-point images in [0, 1]. Pass it explicitly for other ranges, as data_range=target.max() for MRI magnitudes.

psnr 

psnr(pred: Tensor, target: Tensor, data_range: float | None = None) -> float

Return the peak signal-to-noise ratio in dB, pooling all elements.

Parameters:

  • pred (Tensor) –

    Reconstruction.

  • target (Tensor) –

    Reference signal.

  • data_range (float | None, default: None ) –

    Peak value; None uses 1 for floating-point targets and the dtype's maximum for integer ones.

Examples:

>>> round(psnr(torch.zeros(3, 4, 4), torch.full((3, 4, 4), 0.1)), 2)
20.0
Source code in inro/metrics.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def psnr(pred: Tensor, target: Tensor, data_range: float | None = None) -> float:
    """Return the peak signal-to-noise ratio in dB, pooling all elements.

    Args:
        pred: Reconstruction.
        target: Reference signal.
        data_range: Peak value; ``None`` uses ``1`` for floating-point targets and the
            dtype's maximum for integer ones.

    Examples:
        >>> round(psnr(torch.zeros(3, 4, 4), torch.full((3, 4, 4), 0.1)), 2)
        20.0
    """
    return peak_signal_noise_ratio(
        pred.float(), target.float(), data_range=_data_range(target, data_range)
    ).item()

ssim 

ssim(pred: Tensor, target: Tensor, data_range: float | None = None, **kwargs: Any) -> float

Return the structural similarity of two (channels, H, W) images.

Uses an 11x11 Gaussian window by default.

Parameters:

  • pred (Tensor) –

    Reconstruction.

  • target (Tensor) –

    Reference signal.

  • data_range (float | None, default: None ) –

    Peak value; None uses 1 for floating-point targets and the dtype's maximum for integer ones.

  • **kwargs (Any, default: {} ) –

    Passed to structural_similarity_index_measure of torchmetrics.

Source code in inro/metrics.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def ssim(
    pred: Tensor, target: Tensor, data_range: float | None = None, **kwargs: Any
) -> float:
    """Return the structural similarity of two ``(channels, H, W)`` images.

    Uses an 11x11 Gaussian window by default.

    Args:
        pred: Reconstruction.
        target: Reference signal.
        data_range: Peak value; ``None`` uses ``1`` for floating-point targets and the
            dtype's maximum for integer ones.
        **kwargs: Passed to ``structural_similarity_index_measure`` of torchmetrics.
    """
    return structural_similarity_index_measure(
        pred.float().unsqueeze(0),
        target.float().unsqueeze(0),
        data_range=_data_range(target, data_range),
        **kwargs,
    ).item()

ms_ssim 

ms_ssim(pred: Tensor, target: Tensor, data_range: float | None = None, **kwargs: Any) -> float

Return the multi-scale structural similarity of two (channels, H, W) images.

The five default scales require both image sides to be at least 176 pixels.

Parameters:

  • pred (Tensor) –

    Reconstruction.

  • target (Tensor) –

    Reference signal.

  • data_range (float | None, default: None ) –

    Peak value; None uses 1 for floating-point targets and the dtype's maximum for integer ones.

  • **kwargs (Any, default: {} ) –

    Passed to multiscale_structural_similarity_index_measure of torchmetrics.

Source code in inro/metrics.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def ms_ssim(
    pred: Tensor, target: Tensor, data_range: float | None = None, **kwargs: Any
) -> float:
    """Return the multi-scale structural similarity of two ``(channels, H, W)`` images.

    The five default scales require both image sides to be at least 176 pixels.

    Args:
        pred: Reconstruction.
        target: Reference signal.
        data_range: Peak value; ``None`` uses ``1`` for floating-point targets and the
            dtype's maximum for integer ones.
        **kwargs: Passed to ``multiscale_structural_similarity_index_measure`` of
            torchmetrics.
    """
    return multiscale_structural_similarity_index_measure(
        pred.float().unsqueeze(0),
        target.float().unsqueeze(0),
        data_range=_data_range(target, data_range),
        **kwargs,
    ).item()

lpips 

lpips(pred: Tensor, target: Tensor, data_range: float | None = None, net: _LPIPSNet = 'vgg') -> float

Return the learned perceptual image patch similarity of two (3, H, W) images.

Lower is better. Inputs are clipped to the data range. The network, which matches the reference lpips package, is loaded once per backbone and device.

Parameters:

  • pred (Tensor) –

    Reconstruction.

  • target (Tensor) –

    Reference signal.

  • data_range (float | None, default: None ) –

    Peak value; None uses 1 for floating-point targets and the dtype's maximum for integer ones.

  • net (_LPIPSNet, default: 'vgg' ) –

    Backbone of the network.

References

Zhang et al., "The Unreasonable Effectiveness of Deep Features as a Perceptual Metric", CVPR 2018.

Source code in inro/metrics.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
@torch.no_grad()
def lpips(
    pred: Tensor,
    target: Tensor,
    data_range: float | None = None,
    net: _LPIPSNet = "vgg",
) -> float:
    """Return the learned perceptual image patch similarity of two ``(3, H, W)`` images.

    Lower is better. Inputs are clipped to the data range. The network, which matches
    the reference ``lpips`` package, is loaded once per backbone and device.

    Args:
        pred: Reconstruction.
        target: Reference signal.
        data_range: Peak value; ``None`` uses ``1`` for floating-point targets and the
            dtype's maximum for integer ones.
        net: Backbone of the network.

    References:
        Zhang et al., "The Unreasonable Effectiveness of Deep Features as a Perceptual
        Metric", CVPR 2018.
    """
    peak = _data_range(target, data_range)
    pred, target = ((x.float() / peak).clamp(0, 1).unsqueeze(0) for x in (pred, target))
    return _lpips_network(net, pred.device)(pred, target).item()

nmse 

nmse(pred: Tensor, target: Tensor) -> float

Return the normalized MSE ||pred - target||^2 / ||target||^2, as in fastMRI.

Parameters:

  • pred (Tensor) –

    Reconstruction.

  • target (Tensor) –

    Reference signal.

Examples:

>>> nmse(torch.zeros(4), torch.ones(4))
1.0
Source code in inro/metrics.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
def nmse(pred: Tensor, target: Tensor) -> float:
    """Return the normalized MSE ``||pred - target||^2 / ||target||^2``, as in fastMRI.

    Args:
        pred: Reconstruction.
        target: Reference signal.

    Examples:
        >>> nmse(torch.zeros(4), torch.ones(4))
        1.0
    """
    pred, target = pred.float(), target.float()
    return ((pred - target).square().sum() / target.square().sum()).item()

iou 

iou(pred: Tensor, target: Tensor) -> float

Return the intersection over union of occupancies, occupied above zero.

Parameters:

  • pred (Tensor) –

    Reconstruction.

  • target (Tensor) –

    Reference signal.

Examples:

>>> iou(torch.tensor([1.0, 1.0, -1.0]), torch.tensor([1.0, -1.0, -1.0]))
0.5
Source code in inro/metrics.py 
158
159
160
161
162
163
164
165
166
167
168
169
def iou(pred: Tensor, target: Tensor) -> float:
    """Return the intersection over union of occupancies, occupied above zero.

    Args:
        pred: Reconstruction.
        target: Reference signal.

    Examples:
        >>> iou(torch.tensor([1.0, 1.0, -1.0]), torch.tensor([1.0, -1.0, -1.0]))
        0.5
    """
    return binary_jaccard_index((pred > 0).int(), (target > 0).int()).item()

bits_per_pixel 

bits_per_pixel(encoded: bytes, image: Tensor) -> float

Return the encoded size in bits per pixel of a (channels, *spatial) image.

Parameters:

  • encoded (bytes) –

    Encoded signal.

  • image (Tensor) –

    Image whose spatial size counts the pixels.

Examples:

>>> bits_per_pixel(b"123", torch.zeros(3, 4, 6))
1.0
Source code in inro/metrics.py 
172
173
174
175
176
177
178
179
180
181
182
183
def bits_per_pixel(encoded: bytes, image: Tensor) -> float:
    """Return the encoded size in bits per pixel of a ``(channels, *spatial)`` image.

    Args:
        encoded: Encoded signal.
        image: Image whose spatial size counts the pixels.

    Examples:
        >>> bits_per_pixel(b"123", torch.zeros(3, 4, 6))
        1.0
    """
    return len(encoded) * 8 / math.prod(image.shape[1:])

luma 

luma(image: Tensor) -> Tensor

Return the BT.601 luma of a uint8 RGB image, as MATLAB's rgb2ycbcr.

Super-resolution reports PSNR and SSIM on this channel, with data_range=255.

Parameters:

  • image (Tensor) –

    Image (3, H, W).

Returns:

  • Tensor –

    Float image (1, H, W) in [16, 235].

Examples:

>>> white = torch.full((3, 2, 2), 255, dtype=torch.uint8)
>>> round(luma(white).max().item(), 2)
235.0
Source code in inro/metrics.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def luma(image: Tensor) -> Tensor:
    """Return the BT.601 luma of a uint8 RGB image, as MATLAB's ``rgb2ycbcr``.

    Super-resolution reports PSNR and SSIM on this channel, with ``data_range=255``.

    Args:
        image: Image ``(3, H, W)``.

    Returns:
        Float image ``(1, H, W)`` in ``[16, 235]``.

    Examples:
        >>> white = torch.full((3, 2, 2), 255, dtype=torch.uint8)
        >>> round(luma(white).max().item(), 2)
        235.0
    """
    weights = torch.tensor([65.481, 128.553, 24.966], device=image.device) / 255
    return 16 + torch.einsum("c...,c->...", image.float(), weights).unsqueeze(0)

super_resolution_metrics 

super_resolution_metrics(border: int = 0, device: str | device | None = None) -> dict[str, Callable[[Tensor, Tensor], float]]

Return the super-resolution protocol's metrics, as metrics arguments.

These are PSNR and SSIM of the luma channel and LPIPS of RGB. train and evaluate call them as fn(reconstruction, original) on uint8 (3, H, W) images.

Parameters:

  • border (int, default: 0 ) –

    Pixels shaved off each side before scoring, by convention the scale factor.

  • device (str | device | None, default: None ) –

    Device to score on; None keeps the images where they are.

Returns:

  • dict[str, Callable[[Tensor, Tensor], float]] –

    The functions under the keys psnr, ssim and lpips.

Examples:

>>> sorted(super_resolution_metrics(border=4))
['lpips', 'psnr', 'ssim']
Source code in inro/metrics.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
def super_resolution_metrics(
    border: int = 0, device: str | torch.device | None = None
) -> dict[str, Callable[[Tensor, Tensor], float]]:
    """Return the super-resolution protocol's metrics, as ``metrics`` arguments.

    These are PSNR and SSIM of the :func:`luma` channel and LPIPS of RGB.
    :func:`~inro.train` and :func:`~inro.evaluate` call them as
    ``fn(reconstruction, original)`` on uint8 ``(3, H, W)`` images.

    Args:
        border: Pixels shaved off each side before scoring, by convention the scale
            factor.
        device: Device to score on; ``None`` keeps the images where they are.

    Returns:
        The functions under the keys ``psnr``, ``ssim`` and ``lpips``.

    Examples:
        >>> sorted(super_resolution_metrics(border=4))
        ['lpips', 'psnr', 'ssim']
    """

    def shave(image: Tensor) -> Tensor:
        height, width = image.shape[-2:]
        return image.to(device)[:, border : height - border, border : width - border]

    def of_luma(metric: Callable[..., float]) -> Callable[[Tensor, Tensor], float]:
        return lambda pred, target: metric(
            luma(shave(pred)), luma(shave(target)), data_range=255
        )

    return {
        "psnr": of_luma(psnr),
        "ssim": of_luma(ssim),
        "lpips": lambda pred, target: lpips(shave(pred), shave(target)),
    }

Losses

losses

Training losses and differentiable entropy estimates.

A loss follows loss_fn(batch, model) -> (loss, metrics): it runs the model on batch["input"] and returns a scalar loss and logging values, kept as scalar tensors so that train reads them only when it logs.

soft_entropy 

soft_entropy(x: Tensor, num_bins: int | None = None, temperature: float = 1.0, eps: float = 1e-16) -> Tensor

Estimate entropy in bits with differentiable soft bin assignments.

Bin centers span the detached input range. Each element gets a softmax assignment from its squared distance to them, and the entropy is that of the mean assignment. Gradients flow through the assignments, not the bin layout.

Parameters:

  • x (Tensor) –

    Input tensor; all elements are pooled.

  • num_bins (int | None, default: None ) –

    Number of bins; None uses the Freedman–Diaconis width 2 * IQR(x) / N**(1/3), rounding the count up, or Sturges' rule when the IQR vanishes.

  • temperature (float, default: 1.0 ) –

    Assignment kernel width relative to bin width; larger is softer.

  • eps (float, default: 1e-16 ) –

    Constant added inside the logarithm.

Returns:

  • Tensor –

    Scalar entropy, zero for a constant input.

Examples:

>>> x = torch.tensor([-1.0, 1.0])
>>> soft_entropy(x, num_bins=2, temperature=0.1).item()
1.0
Source code in inro/losses.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def soft_entropy(
    x: Tensor,
    num_bins: int | None = None,
    temperature: float = 1.0,
    eps: float = 1e-16,
) -> Tensor:
    """Estimate entropy in bits with differentiable soft bin assignments.

    Bin centers span the detached input range. Each element gets a softmax assignment
    from its squared distance to them, and the entropy is that of the mean assignment.
    Gradients flow through the assignments, not the bin layout.

    Args:
        x: Input tensor; all elements are pooled.
        num_bins: Number of bins; ``None`` uses the Freedman–Diaconis width
            ``2 * IQR(x) / N**(1/3)``, rounding the count up, or Sturges' rule when the
            IQR vanishes.
        temperature: Assignment kernel width relative to bin width; larger is softer.
        eps: Constant added inside the logarithm.

    Returns:
        Scalar entropy, zero for a constant input.

    Examples:
        >>> x = torch.tensor([-1.0, 1.0])
        >>> soft_entropy(x, num_bins=2, temperature=0.1).item()
        1.0
    """
    x_flat = x.reshape(-1, 1)

    x_min = x_flat.detach().min()
    x_max = x_flat.detach().max()

    if x_max == x_min:  # a constant has no entropy
        return (x_flat * 0).sum()

    if num_bins is None:
        iqr_value = _iqr(x_flat.detach())
        num_elements = x_flat.numel()
        bin_width = 2 * iqr_value / (num_elements ** (1 / 3))
        if bin_width > 0:
            num_bins = int(torch.ceil((x_max - x_min) / bin_width).item())
        else:  # no Freedman-Diaconis width: Sturges' rule
            num_bins = math.ceil(math.log2(num_elements)) + 1
            bin_width = (x_max - x_min) / num_bins
    else:
        bin_width = (x_max - x_min) / num_bins

    bin_centers = torch.linspace(x_min, x_max, num_bins, dtype=x.dtype, device=x.device)

    squared_distances = (x_flat - bin_centers) ** 2
    sigma = bin_width * temperature
    logits = -squared_distances / (2 * sigma**2)
    assignments = F.softmax(logits, dim=1)

    probabilities = assignments.mean(dim=0)
    return -(probabilities * torch.log2(probabilities + eps)).sum()

laplace_entropy 

laplace_entropy(x: Tensor, eps: float = 1e-16) -> Tensor

Return the differential entropy of a zero-mean Laplace fit, in bits.

Uses the maximum-likelihood scale b = mean(abs(x)) and the entropy log2(2 * b * e + eps).

Parameters:

  • x (Tensor) –

    Input tensor; all elements are pooled.

  • eps (float, default: 1e-16 ) –

    Constant added inside the logarithm.

Source code in inro/losses.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def laplace_entropy(x: Tensor, eps: float = 1e-16) -> Tensor:
    """Return the differential entropy of a zero-mean Laplace fit, in bits.

    Uses the maximum-likelihood scale ``b = mean(abs(x))`` and the entropy
    ``log2(2 * b * e + eps)``.

    Args:
        x: Input tensor; all elements are pooled.
        eps: Constant added inside the logarithm.
    """
    scale = x.abs().mean()
    return torch.log2(2 * scale * math.e + eps)

rate_distortion_loss 

rate_distortion_loss(batch: dict[str, Any], model: Module, alpha: float = 0.0, eps: float = 1e-16) -> tuple[Tensor, dict[str, Tensor]]

Balance reconstruction MSE with estimated parameter entropy.

The loss is MSE + alpha * entropy, the entropy being the size-weighted mean of laplace_entropy over the trainable parameters, evaluated only when alpha > 0.

Parameters:

  • batch (dict[str, Any]) –

    Batch with input and target.

  • model (Module) –

    INR to evaluate.

  • alpha (float, default: 0.0 ) –

    Weight of the rate term.

  • eps (float, default: 1e-16 ) –

    Constant added inside the logarithm.

Returns:

  • tuple[Tensor, dict[str, Tensor]] –

    Scalar loss and the metrics mse, psnr (for targets in [-1, 1]), loss, and entropy when alpha > 0, as scalar tensors.

Source code in inro/losses.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def rate_distortion_loss(
    batch: dict[str, Any],
    model: nn.Module,
    alpha: float = 0.0,
    eps: float = 1e-16,
) -> tuple[Tensor, dict[str, Tensor]]:
    """Balance reconstruction MSE with estimated parameter entropy.

    The loss is ``MSE + alpha * entropy``, the entropy being the size-weighted mean of
    :func:`laplace_entropy` over the trainable parameters, evaluated only when
    ``alpha > 0``.

    Args:
        batch: Batch with ``input`` and ``target``.
        model: INR to evaluate.
        alpha: Weight of the rate term.
        eps: Constant added inside the logarithm.

    Returns:
        Scalar loss and the metrics ``mse``, ``psnr`` (for targets in ``[-1, 1]``),
            ``loss``, and ``entropy`` when ``alpha > 0``, as scalar tensors.
    """
    output = model(batch["input"])
    target = batch["target"]
    metrics: dict[str, Tensor] = {}

    mse = F.mse_loss(output, target)
    metrics["mse"] = mse.detach()
    metrics["psnr"] = 10 * torch.log10(4.0 / metrics["mse"])

    if alpha > 0:
        entropy = 0
        total_params = 0
        for param in model.parameters():
            if not param.requires_grad:
                continue
            num_params = param.numel()
            total_params += num_params
            entropy += num_params * laplace_entropy(param, eps=eps)

        entropy /= total_params
        loss = mse + alpha * entropy
        metrics["entropy"] = entropy.detach()
    else:
        loss = mse

    metrics["loss"] = loss.detach()
    return loss, metrics

sdf_loss 

sdf_loss(batch: dict[str, Any], model: Module) -> tuple[Tensor, dict[str, Tensor]]

Return the L1 signed-distance loss and the metrics loss and l1.

Parameters:

  • batch (dict[str, Any]) –

    Batch with input and target.

  • model (Module) –

    INR to evaluate.

Source code in inro/losses.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def sdf_loss(
    batch: dict[str, Any], model: nn.Module
) -> tuple[Tensor, dict[str, Tensor]]:
    """Return the L1 signed-distance loss and the metrics ``loss`` and ``l1``.

    Args:
        batch: Batch with ``input`` and ``target``.
        model: INR to evaluate.
    """
    output = model(batch["input"])
    target = batch["target"]

    loss = F.l1_loss(output, target)
    value = loss.detach()
    return loss, {"loss": value, "l1": value}

super_resolution_loss 

super_resolution_loss(batch: dict[str, Any], model: Module) -> tuple[Tensor, dict[str, Tensor]]

Return the L1 loss through the degradation and the metrics l1 and psnr.

Renders each block's footprint, downsamples it with the block's kernel rows along both axes and matches the low-resolution pixels. The PSNR is for targets in [-1, 1].

Parameters:

Source code in inro/losses.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def super_resolution_loss(
    batch: dict[str, Any], model: nn.Module
) -> tuple[Tensor, dict[str, Tensor]]:
    """Return the L1 loss through the degradation and the metrics ``l1`` and ``psnr``.

    Renders each block's footprint, downsamples it with the block's kernel rows along
    both axes and matches the low-resolution pixels. The PSNR is for targets in
    ``[-1, 1]``.

    Args:
        batch: Batch of :class:`~inro.SuperResolutionDataset` items.
        model: INR to evaluate.
    """
    weights = batch["weight_y"], batch["weight_x"]
    output = downsample(model(batch["input"]), weights)
    target = batch["target"]

    loss = F.l1_loss(output, target)
    mse = F.mse_loss(output.detach(), target)
    return loss, {"l1": loss.detach(), "psnr": 10 * torch.log10(4.0 / mse)}