SR-Forge

Metrics

Metrics and losses.

A metric measures how good an output is. A loss is a metric you train on: the same class, with a weight that decides whether it drives the gradient. Modules are grouped by what a metric assumes about the data:

• :mod:~srforge.metrics.regression — any tensor: L1, MSE, Charbonnier, PSNR
• :mod:~srforge.metrics.classification — yes/no per element: F1, IoU, MCC, counts; Dice / Jaccard / Tversky losses, binary and multi-class cross-entropy
• :mod:~srforge.metrics.image — images: SSIM, MGE, TotalVariation, TBE, LPIPS, SAM
• :mod:~srforge.metrics.wrappers — change how another metric is evaluated
• :mod:~srforge.metrics.regularization — penalties on a single tensor
• :mod:~srforge.metrics.adversarial — GAN objectives
• :mod:~srforge.metrics.adapters — torchmetrics and pyiqa

The core API is importable from here directly.

Metric

Bases: Module, IOModule, ABC

Base class for metrics — and so for losses, which are metrics you train on.

Subclasses pick ONE of two implementation levels:

• Level 2 (compositional, preferred for new metrics). Override :meth:pointwise only — return a per-element score map with the same spatial shape as the inputs. The default :meth:calculate_score composes :meth:pointwise → :meth:reduce (with optional masking applied inside reduce). Free for wrapping by modifiers like :class:UncertaintyLoss.

• Level 1 (structural, full control). Override :meth:calculate_score directly to return a [B] tensor. Use this when the metric has no per-pixel value (SSIM windows, MGE gradients), or when something non-linear follows the average (PSNR's logarithm). Skip :meth:pointwise — the default raises a clear :class:NotImplementedError, which is exactly what should happen if someone tries to wrap a structural metric with :class:UncertaintyLoss.

Subclasses with non-canonical signatures (e.g. :class:UncertaintyLoss takes (x, y, sigma)) must override :meth:calculate_score as well to route their extra inputs.

Parameters:

Name	Type	Description	Default
weight	float	Scalar weight applied during LossCombiner and MetricScores.total_weighted. Use 0 to track a metric without including it in backward.	1.0
name	str	Human-readable name for logging. Defaults to the class name.	None
reduction	str	How reduce() collapses the pointwise score map to [B]. One of: "mean" (default) — average over non-batch dims. With a mask, divides by the unmasked pixel count. "sum" — sum over non-batch dims. Useful for proper Gaussian-NLL math where the inner term should be summed, not averaged. "none" — return the score map unchanged; the caller handles reduction. Only a metric whose score goes through :meth:reduce has a map to collapse — every Level-2 metric, and a Level-1 metric that calls self.reduce(...) and says so with uses_reduction = True (PSNR, Regularizer). Any other metric refuses "sum" and "none" instead of ignoring them.	'mean'
aggregate	Optional[str]	How per-image values combine — over a training batch for the value you backpropagate, and over the epoch for the value you report. Each class states what it supports; the first is its default, used when this is left None. "macro" — average the per-image values. Every image counts equally. The default for most metrics. "micro" — pool, then compute once. For a pointwise metric, the mean over every valid pixel, so large and mostly-unmasked images count for more. Available automatically on any Level-2 metric that uses the default calculate_score (with reduction="mean"), and on metrics that define a pooled form (PSNR, SAM, the classification ratios, the torchmetrics adapter). "custom" — the class overrides :meth:finalize with its own epoch rule (a total, a worst case). Assigned automatically, never chosen. Anything a class does not support is refused here, naming what it does support.	None

reductions property

How each part of this metric's value combines over the epoch.

None (the default) keeps every per-sample value, which is what every metric that predates this does. Declare otherwise when a part is not something you average::

@property
def reductions(self):
    return {"tp": "sum", "fp": "sum", "fn": "sum"}

Valid rules are sum, mean, min, max and cat. Either a mapping (for a metric returning a dict of parts) or a single string (for one returning a plain tensor).

This exists so :meth:finalize receives values already combined according to what they are, instead of having to survive an aggregation it did not choose. Without it, a total could not be reported at all — a distributed run averages every key, so a sum comes back divided by the number of images — and neither could a worst case. Even a ratio was fragile: adding the +1 smoothing that segmentation code uses everywhere silently changed the answer.

__init_subclass__(**kwargs)

Install the Metric-argument check on any subclass that needs it, and check what the class says about aggregation.

Preserved and invoked by audit_subclasses, which wraps the base and calls whatever __init_subclass__ the class already defines.

set_io(io_cfg, *, strict=True, require_all=True)

Set IO binding using the unified format: {"inputs": {param: field}}.

Metrics only read from Entry (no outputs). If outputs is provided, a :class:TypeError is raised.

finalize(raw: Any) -> torch.Tensor

Turn this metric's accumulated values into the number to report.

The third and last stage, after :meth:pointwise (per pixel) and :meth:calculate_score (per image). The default is what every metric in SR-Forge does today: average the per-image scores collected over the epoch.

Override it when that average is the wrong answer. A micro-averaged F1 is 2*ΣTP / (2*ΣTP + ΣFP + ΣFN) over every pixel seen, which is not the mean of per-image F1 values — on change-detection data the two differ by tens of percent, because images containing no change score 0 and drag a mean down. A pooled PSNR is the log of the pooled MSE, not the mean of per-image PSNRs; measured on SR data the gap is over 10 dB.

A metric that overrides this returns its ingredients rather than a finished score — {"tp": [B], "fp": [B], "fn": [B]}, say — and combines them here. The aggregation then travels with the metric that owns it, instead of living somewhere that has to find it again later.

Store what adds up. Counts and sums, never means or ratios: a mean has already discarded the weight it was taken over, so averaging means silently mis-weights images of different sizes. If the quantity has a denominator, carry the numerator and denominator separately. A :meth:pointwise metric gets this for free — see below.

That rule also makes the result correct under DDP, where reduce_scores collapses each key to a global mean before this runs: a ratio of sums is unchanged by that, a bare sum() is not. An order statistic (median, percentile) cannot be expressed here at all, because the individual values no longer exist.

Under aggregate="micro" on a pointwise metric the default does the pooling itself: raw is {"sum", "count"} already totalled, and the value is their ratio — the mean over every valid pixel seen.

calculate_score(*args, **kwargs) -> torch.Tensor

Compute the per-sample value.

The base method is declared as *args, **kwargs so subclasses can override it with whatever inputs they need (IO binding routes Entry fields to the declared params — e.g. ShiftNetRegisLoss.calculate_score(shifts)) without tripping static-analysis Liskov checks. The polymorphic arity is intentional.

Subclasses that don't override this method inherit the default Level-2 path (:meth:pointwise → mask → :meth:reduce), implemented in :meth:_default_calculate_score. Their inputs are :meth:pointwise's parameters plus y_mask.

Override directly for Level-1 (structural) metrics that aren't decomposable into a per-element operation (SSIM, MGE), whose per-image score is a non-linear function of a mean (PSNR), or with a non-canonical input signature.

pointwise(*args, **kwargs) -> torch.Tensor

Per-pixel score map — no masking, no reduction.

Its parameters are the metric's inputs (plus y_mask, which the framework applies): pointwise(self, x) for a no-reference metric, pointwise(self, x, y) for the usual comparison, UncertaintyLoss.pointwise(self, x, y, log_var) for more. Name and annotate them like calculate_score parameters; y_mask itself is reserved. The base method is declared as *args, **kwargs so subclasses can choose freely without tripping static-analysis Liskov checks.

Override in subclasses that decompose into a per-element comparison (MSE, L1, Charbonnier, CrossEntropy, SAM, …). The default raises so wrapping a non-pointwise metric with a Level-2 modifier (e.g. :class:UncertaintyLoss) fails loudly instead of silently producing garbage.

Contract:

• Batch dim first; the output's first axis is [B].

• Non-batch dims may differ from the inputs, as long as :meth:reduce can collapse them to a single per-sample scalar (it does view(B, -1) internally). Typical shapes:

• [B, C, H, W] — element-wise comparison preserving all input dims (MSE, L1, Charbonnier).

• [B, 1, H, W] — comparison that collapses the channel axis (SAM's per-pixel spectral angle).

What does NOT fit cleanly: shapes that drop spatial pixels (windowed reductions like SSIM, forward-difference like TotalVariation) — those stay Level-1 (override :meth:calculate_score directly). * No masking — the default :meth:calculate_score dispatches masking to :meth:reduce (which zeros the score map at masked positions and normalises the denominator). Pointwise sees raw inputs and may produce arbitrary finite values at "masked" positions — they're discarded by reduce. If your pointwise is convolutional / windowed (SSIM, MGE) and needs input neutralisation so windows never reach masked pixels, use :meth:erode_mask in your own :meth:calculate_score override. * No reduction — :meth:reduce does that, separately and configurably.

No :meth:calculate_score override is needed for extra inputs: the inherited one passes every input through to pointwise.

reduce(score_map: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor

Reduce a [B, ...] per-element score map to [B].

With a mask: brings it to the map's shape (see :meth:_match_mask), zeros masked positions, sums the remaining, and normalises by the number of valid elements (clamped to ≥ 1 so all-masked samples return 0, not NaN). Without a mask: averages or sums over all non-batch dims, depending on :attr:_reduction.

Returns the score map unchanged when reduction == "none".

__check_annotated_signatures()

Ensures that the calculate_score method has annotated signatures for all parameters.

erode_mask(mask: torch.Tensor, window_size: int) -> torch.Tensor staticmethod

Keep only the positions whose whole window_size window is valid.

For windowed metrics (SSIM, gradient metrics, anything built on a convolution). The value such a metric computes at a pixel depends on every pixel in the window around it, so a window that reaches into a masked area compares values nobody vouches for. Zeroing the masked inputs first does not help — the zeros become part of the windows next to them: SSIM then sees a perfect match, a gradient filter sees an edge that is not in the image.

Compute the metric's map on the real inputs, then average it over this eroded mask. Each value is the smallest mask value in its window (a min-filter), so soft masks stay soft. Positions outside the image are treated as valid: image borders are the metric's own padding business, not the mask's, and an all-ones mask therefore changes nothing.

mask_pixels(sr: torch.Tensor, hr: torch.Tensor, hr_mask=None) staticmethod

Neutralise inputs at masked positions: returns (sr*mask, hr*mask, unmasked_count_clamped_to_one).

Not for windowed metrics. Zeroing masked positions before a window runs puts those zeros into the neighbouring windows — SSIM and the built-in gradient metric both scored masked areas wrongly this way. Use :meth:erode_mask on the metric's map instead. Pointwise metrics don't need this either — :meth:reduce already zeros the score map at masked positions and normalises the denominator.

The unmasked_count is clamped to ≥1 so callers that divide by it don't hit div-by-zero on samples whose mask is entirely zero. With sr*mask = hr*mask = 0 in that case, the resulting per-sample value is exactly 0 instead of NaN — what we want for label-less / placeholder HR samples (challenge_dev mixed into training, etc.). Pixel counts are integers ≥ 1 in the normal partial-mask case so the clamp is a no-op there.

MetricEntry dataclass

Holds all variants of a single metric

• raw: output of calculate_score (no sign, no weight)
• best_min: True if smaller is better; False if larger is better
• weight: optimization weight (can be 0)
• finalize: optional rule for turning accumulated raw into one number, replacing the default "average the per-sample values"

finalize exists because some metrics are not averages of per-sample scores. A micro-averaged F1 is 2*ΣTP / (2*ΣTP + ΣFP + ΣFN) over every pixel in the epoch, which is not the mean of per-image F1 — on change-detection data the two differ by tens of percent, because images containing no change score 0 and drag a mean down. Such a metric reports its parts ({"tp": [B], "fp": [B], "fn": [B]}) and supplies the rule that combines them, so the aggregation travels with the value instead of living in something that has to find it again by name later.

A finalized entry is not a per-sample score, so it never appears in the per-sample views. What it contributes to total_raw and total_weighted depends on aggregate:

• "micro" — its pooled value over whatever the entry currently holds: one training batch when the runner backpropagates, the whole epoch when a checkpoint compares totals. That is batch-level micro for training, and it is differentiable wherever the metric is.
• "custom" — nothing. A count or a worst case is not a loss term.

The rule must be a function of per-key means. In a single process raw holds every per-sample value; in a distributed run TrainingStrategy.reduce_scores has already collapsed each key to one global mean. A rule written with .mean(), or as a ratio of sums (where the sample count cancels), gives the same answer in both cases. F1, precision, recall, IoU, MCC and a pooled PSNR all satisfy this.

A rule that does not — a bare .sum(), or an order statistic such as a median or a 95^th percentile — is correct on one process and silently wrong on N, because the values it needs no longer exist by the time it runs. Such a metric does not belong here.

is_finalized: bool property

Whether this entry carries parts + a rule rather than scores.

counts_toward_total: bool property

Whether this entry is a loss term. Per-image scores are; a pooled (micro) value is the same quantity weighted differently, so it is too; a custom epoch rule (a total, a worst case) is not.

total_batch(weighted: bool) -> torch.Tensor

This entry's [B] contribution to a loss total.

A pooled entry has one value for the whole batch rather than one per image, so it is repeated across the batch: the runner's .mean() then recovers exactly that value, and its gradient, with no rescaling.

reduction_for(key: str = None) -> str

The declared reduction for key, defaulting to "cat".

reduced() -> TensorOrDict

raw with each key collapsed as the metric declared.

This is what :meth:finalize receives: values already totalled, minimised or averaged according to what they are, so the formula does not have to compensate for an aggregation it did not choose. Without a declaration the raw values pass through unchanged, which is the behaviour of every metric that predates this. With bands, each band is collapsed on its own: {band: reduced value}.

band_values(weighted: bool = False) -> Dict[str, torch.Tensor]

Each band's own value, by the metric's own rule.

For a finalized entry with bands: the metric's :meth:finalize run on each band's totals, so F1.red is the micro F1 of the red band alone. weighted signs and weights micro values the way :meth:final_weighted does.

final_value() -> torch.Tensor

The single number this entry represents.

With bands, the bands combine the way the metric's images do:

• macro and micro — the average of the band values, each band counting equally (bands often differ in resolution, and pooling would let the finest one dominate);
• custom — the band totals combined by the metric's own declared rules, then its finalize (a count sums, a worst case takes the minimum). Where a rule cannot run across bands, the average of the band values.

final_weighted() -> torch.Tensor

:meth:final_value as it enters the weighted total.

Sign and weight apply to a pooled (micro) value exactly as to a per-image score, so the weighted numbers of a MetricScores add up to its total_weighted(). A custom rule's value — a count, a worst case — is not a loss term and has no weighted form, so it is returned as is.

MetricScores dataclass

__setitem__(name: str, entry: MetricEntry) -> None

In-place add or replace a metric. Used by hooks that inject auxiliary loss terms from on_post_forward — drop a :class:MetricEntry here and the runner's total_weighted() picks it up before backward, and LossLogger emits it as a per-name series alongside the criterion's metrics.

add_scores(other: MetricScores) -> None

In-place accumulation: concatenates batch dimension when names match.

LossScheduler

Epoch-aware loss wrapper that swaps the active Metric at configured milestones.

Parameters:

Name	Type	Description	Default
schedule	Mapping[int, Metric]	Mapping of {epoch_milestone: Metric}. Must contain key 0.	required
start_epoch	int	Current/resume epoch (defaults to 0).	0

__call__(*args, **kwargs) -> MetricScores

Delegate to the currently active loss function.

update(epoch: int) -> LossScheduler

Advance the scheduler to epoch, updating the active loss.

get_loss_fn() -> Metric

Return the currently active Metric instance.

adapt(obj: Any, **options: Any) -> Any

Return obj as a :class:Metric, wrapping it if it comes from a library SR-Forge knows how to adapt.

Anything that is already a Metric — and anything no adapter claims — is returned unchanged and identical, so this is safe to call on every object the resolver builds.

Parameters:

Name	Type	Description	Default
obj	Any	The object to adapt.	required
**options	Any	Forwarded to the adapter's constructor (weight, name, aggregate). Passing options for an object that needs no adapting is an error, because they would be silently dropped.	{}