SR-Forge

Loss

Loss

Bases: Module, IOModule, ABC

Base class for losses.

Subclasses pick ONE of two implementation levels:

• Level 2 (compositional, preferred for new losses). 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 loss isn't decomposable into a per-element operation (SSIM windows, LPIPS feature network, SAM spectral angles). Skip :meth:pointwise — the default raises a clear :class:NotImplementedError, which is exactly what should happen if someone tries to wrap a structural loss 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		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. Ignored when the subclass overrides calculate_score directly (Level 1).	'mean'

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

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

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

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

Compute the per-sample loss/metric value.

Canonical signature: calculate_score(x, y, y_mask=None). The base method is declared as *args, **kwargs so subclasses with non-canonical signatures (IO binding routes arbitrary Entry fields to declared params — e.g. ShiftNetRegisLoss.calculate_score(shifts), UncertaintyLoss.calculate_score(x, y, sigma, y_mask=None)) can override freely without tripping static-analysis Liskov checks. The polymorphic arity is intentional.

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

Override directly for Level-1 (structural) losses that aren't decomposable into a per-element operation (SSIM, LPIPS, SAM), or for losses with a non-canonical input signature.

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

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

Canonical signature: pointwise(x: Tensor, y: Tensor) -> Tensor. The base method is declared as *args, **kwargs so subclasses can extend with extra positional inputs (e.g. UncertaintyLoss.pointwise(x, y, sigma)) without tripping static-analysis Liskov checks — the polymorphic arity is intentional.

Override in subclasses that decompose into a per-element comparison (MSE, L1, Charbonnier, CrossEntropy, SAM, …). The default raises so wrapping a non-pointwise loss 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 before the windowed op, use :meth:mask_pixels in your own :meth:calculate_score override. * No reduction — :meth:reduce does that, separately and configurably.

Implementations with extra positional args (like UncertaintyLoss.pointwise(x, y, sigma)) must also override :meth:calculate_score to thread the extra args through.

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

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

With a mask: zeros masked positions, sums the remaining, and normalises by the unmasked pixel count (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.

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).

Only useful for convolutional / windowed losses (SSIM, MGE, …) that need masked positions zeroed before a windowed op so neighbouring statistics aren't contaminated. Pointwise losses don't need this — :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 loss 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.

cLossCombiner

Bases: LossCombiner

Corrected-loss combiner: - For each underlying loss, evaluates it on a grid of shifted HR patches. - For each pixel, picks min or max over shifts (depending on best_min). - Returns MetricScores with one MetricEntry per underlying loss, preserving each loss's weight and best_min.

.. deprecated:: 0.15.0 Scheduled for removal in v0.16.0. Use :class:~srforge.loss.metrics.CorrectedLoss instead — same math, applied to a single base loss, composable via the framework's standard IO binding. For N corrected metrics, declare N CorrectedLoss instances in your combiner.

forward(entry: Entry) -> MetricScores

Build kwargs for our own calculate_score (sr, hr, hr_mask), handle multi-band dict inputs, and wrap results into MetricScores.

calculate_score(x: torch.Tensor, y: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, torch.Tensor]

Compute raw per-loss scores after spatial correction.

IMPORTANT: - Returns raw scores (no sign flip, no weighting). - Weight and best_min are handled by MetricEntry / MetricScores.