SR-Forge

Losses¶

A Loss evaluates how well your model's predictions match the target. In SR-Forge, losses are components that read from Entry fields (like sr and hr) and return MetricScores — they never write back to the Entry. This makes them read-only consumers of the pipeline output.

Dataset --> Transforms --> Model --> [Entry with sr, hr, ...] --> Loss --> MetricScores

Your First Loss¶

from srforge.loss import Loss
import torch

class MyL1(Loss):
    @property
    def best_min(self) -> bool:
        return True  # lower is better

    def calculate_score(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return (x - y).abs().mean(dim=(-3, -2, -1))  # per-sample scalar [B]

Three things to notice:

best_min — tells the framework whether to minimize (True) or maximize (False) this metric during optimization.
calculate_score() — your computation. Parameter names are used for IO binding (here: x and y).
Type annotations are required — every parameter in calculate_score() must be annotated. This drives the automatic multi-band splitting behavior.

IO Binding¶

Losses use the same set_io() format as every other SR-Forge component, but with inputs only — losses don't produce outputs.

from srforge.loss.metrics import L1
from srforge.data import Entry

loss = L1()
loss.set_io({"inputs": {"x": "sr", "y": "hr"}})

entry = Entry(sr=torch.randn(1, 3, 64, 64), hr=torch.randn(1, 3, 64, 64))
scores = loss(entry)  # reads entry["sr"] → x, entry["hr"] → y

Identity default — when set_io() is not called, the loss maps parameter names to same-named Entry fields. If your Entry has fields x and y, you don't need set_io() at all:

loss = L1()
entry = Entry(x=torch.randn(1, 3, 64, 64), y=torch.randn(1, 3, 64, 64))
scores = loss(entry)  # x → x, y → y — just works

No outputs for losses

Losses return MetricScores, not Entry fields. Passing outputs to set_io() raises TypeError.

Calling Modes¶

Losses support three calling modes:

1. Entry-based (primary)¶

loss = L1()
loss.set_io({"inputs": {"x": "prediction", "y": "ground_truth"}})
scores = loss(entry)

2. Positional tensors¶

loss = L1()
scores = loss(pred_tensor, target_tensor)  # maps to (x, y) in signature order

3. Keyword tensors¶

loss = L1()
scores = loss(x=pred_tensor, y=target_tensor)

Note

When using keyword mode, keys must match the Entry field names (after mapping), not the parameter names. If you set set_io({"inputs": {"x": "sr", "y": "hr"}}), use loss(sr=..., hr=...).

Masking¶

Most built-in losses accept an optional y_mask parameter for spatial masking — useful when parts of the target image are invalid (e.g., borders, clouds, no-data regions).

loss = L1()
loss.set_io({"inputs": {"x": "sr", "y": "hr", "y_mask": "mask"}})

entry = Entry(
    sr=torch.randn(1, 3, 64, 64),
    hr=torch.randn(1, 3, 64, 64),
    mask=torch.ones(1, 3, 64, 64),  # 1 = valid, 0 = ignore
)
scores = loss(entry)

For pointwise losses, masking is handled automatically — reduce() zeros the score map at masked positions and normalises by the unmasked pixel count (clamped to ≥1, so all-masked samples return 0, not NaN — useful when challenge-dev placeholders are mixed into training). Just override pointwise():

class MyAbsDiff(Loss):
    @property
    def best_min(self):
        return True

    def pointwise(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return (y - x).abs()

For convolutional / windowed losses (SSIM, MGE) that need masked positions zeroed before a windowed op so neighbouring statistics aren't contaminated, the mask_pixels() helper is available:

class MaskedConvLoss(Loss):
    @property
    def best_min(self):
        return True

    def calculate_score(self, x: torch.Tensor, y: torch.Tensor, y_mask: torch.Tensor = None):
        x, y, total_unmasked = self.mask_pixels(x, y, y_mask)
        # ... windowed op on neutralised x, y ...
        return result / total_unmasked

Multi-Band (Dict) Inputs¶

Loss inputs can be tensors or dicts of tensors (one per spectral band). The behavior depends on type annotations:

Tensor-annotated parameters (default) — automatic per-band splitting¶

When calculate_score parameters are annotated as torch.Tensor, dict inputs are automatically split per band. The loss runs once per band, and results are collected into a {band: score} dict:

class BandL1(Loss):
    @property
    def best_min(self):
        return True

    def calculate_score(self, x: torch.Tensor, y: torch.Tensor):
        return (x - y).abs().mean(dim=(-3, -2, -1))

loss = BandL1()
entry = Entry(
    x={"b1": torch.randn(1, 1, 64, 64), "b2": torch.randn(1, 1, 64, 64)},
    y={"b1": torch.randn(1, 1, 64, 64), "b2": torch.randn(1, 1, 64, 64)},
)
scores = loss(entry)
# scores.as_raw_dict() → {"BandL1": {"b1": tensor([...]), "b2": tensor([...])}}

Dict-annotated parameters — single call, full dict passed¶

When parameters are annotated as dict, the entire dict is passed as-is. No splitting occurs — your code handles the dict structure:

class DictSAM(Loss):
    @property
    def best_min(self):
        return True

    def calculate_score(
        self,
        x: dict[str, torch.Tensor],
        y: dict[str, torch.Tensor],
    ):
        # Full dict available — compute spectral angle across bands
        ...

MetricScores¶

Loss.__call__ always returns MetricScores, never raw tensors. The object holds one MetricEntry per loss name, each carrying the raw score plus its weight and best_min orientation. Per-entry shape is Tensor[B] for scalar losses and dict[str, Tensor[B]] for multi-band losses.

Aggregation (the backprop chain)¶

scores = loss(entry)              # MetricScores
total  = scores.total_weighted()  # Tensor[B] — sum across all entries,
                                  # with sign + weight applied
loss_val = total.mean()           # Tensor[] — reduce over batch
loss_val.backward()

The runner does exactly this on every training step (srforge/training/runners.py:240, 244-245 — total_weighted over accumulation steps, divided by gradient_accumulation_steps, then scaler.scale(...).mean().backward()).

Method reference¶

Method	Returns	Purpose
`as_raw_dict()`	`{name: tensor\\|dict}`	Raw scores, no sign flip or weighting
`as_weighted_dict()`	`{name: tensor\\|dict}`	With optimization sign (`+1` for `best_min`, `-1` for `best_max`) + weight applied
`as_raw_flat_dict()`	`{name: tensor}`	Flattened — dict-valued entries become `{name}/{band}`
`as_weighted_flat_dict()`	`{name: tensor}`	Same flattening on the weighted side
`as_summary_dict()`	`{name: tensor\\|dict}`	For logging — uses each entry's `as_opt`
`total_raw()`	`Tensor[B]`	Sum of all raw scores across entries
`total_weighted()`	`Tensor[B]`	Sum with sign + weight (the value the trainer backprops on)
`mean_raw()`	`{name: tensor}`	Per-entry batch mean of raw scores
`mean_weighted()`	`{name: tensor}`	Per-entry batch mean with sign + weight
`merge(other)`	`MetricScores`	Concatenate batches — used by runners to accumulate within an epoch
`add_scores(other)`	`None`	In-place version of `merge`
`detached(cpu=True)`	`MetricScores`	Detached copy for logging without breaking the graph
`__iter__()`	iter[str]	Iterate over entry names
`__getitem__(name)`	`MetricEntry`	Look up an entry by name

MetricEntry¶

Each MetricScores[name] returns a MetricEntry:

Method	Returns	Purpose
`as_raw()`	`tensor\\|dict`	Raw per-sample value
`as_weighted()`	`tensor\\|dict`	Sign + weight applied
`as_opt()`	`tensor\\|dict`	"Lower is better" view (negates `best_max` losses)
`raw_batch()` / `opt_batch()` / `weighted_batch()`	`Tensor[B]`	Same views collapsed to a flat batch tensor

All of these are methods — remember the trailing (). MetricScores.total_weighted() (above) is also a method.

MetricEntry.concat(other) produces a single entry covering both batches; MetricScores.merge calls it on every shared name.

Built-In Losses¶

Loss	Parameters	`best_min`	Description
L1	`x, y, y_mask=None`	min	Mean absolute error (MAE)
MSE	`x, y, y_mask=None`	min	Mean squared error
PSNR	`x, y, y_mask=None`	max	Peak signal-to-noise ratio
Charbonnier	`x, y, y_mask=None`	min	Smooth L1 (epsilon-controlled)
SSIM	`x, y, y_mask=None`	max	Structural similarity index
TotalVariation	`x, y=None, mask=None`	min	Spatial smoothness penalty
MGE	`x, y, y_mask=None`	min	Mean Gradient Error (Sobel-based)
SAM	`x, y`	min	Spectral Angle Mapper
LPIPS	`x, y`	min	Learned Perceptual Image Patch Similarity
CrossEntropy	`x, y`	min	Per-sample mean cross-entropy
TBE	`sr, entry`	min	The Blur Effect (non-reference)
UncertaintyLoss	`base_metric, x, y, log_var, y_mask=None`	min	NLL data term for any pointwise base. Pair with `Regularizer` for the `+log_var` regularizer
Regularizer	`penalty, x, y_mask=None`	min	Generic single-tensor magnitude penalty (identity, abs, square, exp, entropy, huber, …)
CorrectedLoss	`base_metric, x, y, y_mask=None`	inherits base	Shift- and photometric-correction wrapper

All losses accept weight and name constructor parameters.

Deprecated for removal in v0.16.0

cLossCombiner is deprecated. Use CorrectedLoss instead — one wrapper per base metric you want shift-corrected, composes via the standard IO binding.

LossCombiner¶

Combine multiple losses, each with its own weight:

from srforge.loss import L1, SSIM, LossCombiner

combiner = LossCombiner([
    L1(weight=1.0, name="L1").set_io({"inputs": {"x": "sr", "y": "hr"}}),
    SSIM(weight=0.01, name="SSIM").set_io({"inputs": {"x": "sr", "y": "hr"}}),
])

scores = combiner(entry)
# scores.metrics has both "L1" and "SSIM" entries
total = scores.total_weighted()  # weighted sum for backprop

Each loss evaluates independently on the same Entry. The combiner merges all MetricScores — names must be unique across losses.

cLossCombiner¶

A variant that adds spatial correction — it evaluates each loss over a grid of pixel shifts and picks the best alignment per sample. Useful for super-resolution where sub-pixel misalignment between prediction and target is common:

from srforge.loss import L1, cLossCombiner

combiner = cLossCombiner(
    losses=[L1(weight=1.0)],
    border=3,           # search ±3 pixels
    do_correction=True, # apply brightness correction
)
combiner.set_io({"inputs": {"x": "sr", "y": "hr"}})

Loss names are prefixed with c (e.g., L1 becomes cL1).

CorrectedLoss — the modern wrapper analogue¶

CorrectedLoss applies the same shift+photometric correction as cLossCombiner, but to a single base loss instead of combining many under one shift loop. Two reasons to prefer it for new code:

Composes via standard IO binding. It calls self.base(new_entry) internally, so the base loss's own IO binding handles parameter routing. You get the framework's normal "outer io + inner io" pattern without any special wrapper logic.
Simpler config. Each corrected metric is one wrapper instance. To run multiple corrected metrics, just declare multiple CorrectedLoss entries — no combiner needed.

from srforge.loss.metrics import L1, CorrectedLoss

loss = CorrectedLoss(base=L1(), border=3, do_correction=True)
loss.set_io({"inputs": {"x": "sr", "y": "hr", "y_mask": "valid_mask"}})

YAML:

loss:
  _target: srforge.loss.metrics.CorrectedLoss
  params:
    base_metric:
      _target: srforge.loss.metrics.L1
      params: {}
    border: 3
    do_correction: true
  io:
    inputs: {x: sr, y: hr, y_mask: valid_mask}

If the base loss uses non-canonical parameter names (e.g. a custom loss with (sr, hr, mask) instead of (x, y, y_mask)), add an inner io block that maps its params to CorrectedLoss's canonical inner field names (x, y, y_mask):

loss:
  _target: srforge.loss.metrics.CorrectedLoss
  params:
    base_metric:
      _target: my_module.MyLoss            # params: (sr, hr, mask)
      io:
        inputs: {sr: x, hr: y, mask: y_mask}   # ← inner io routes params
    border: 3
  io:
    inputs: {x: sr_field, y: hr_field, y_mask: valid_mask}

cLossCombiner stays available for backward compatibility — but CorrectedLoss is the recommended pattern for new code.

UncertaintyLoss + Regularizer — heteroscedastic NLL¶

The full heteroscedastic-NLL objective is

loss = base(x, y) · exp(-σ) + σ
       └─────────────────┘   └─┘
        data term            regularizer

These are two separate losses in sr-forge:

UncertaintyLoss(base_metric=...) — produces the data term: base.pointwise(x, y) · exp(-σ). Wraps any pointwise loss (MSE / L1 / Charbonnier / …); the base is used only for its pointwise() method.
Regularizer(penalty="identity") — produces the +σ regularizer: reduce(σ). Generic single-tensor magnitude penalty — log-variance regularization is just one application (see Regularizer below for other uses).

Splitting these two responsibilities makes multi-task with shared uncertainty work correctly by construction — no double-counting.

Single-task usage¶

from srforge.loss import LossCombiner
from srforge.loss.metrics import MSE, UncertaintyLoss, Regularizer

loss = LossCombiner([
    UncertaintyLoss(base_metric=MSE()).set_io(
        {"inputs": {"x": "sr", "y": "hr", "log_var": "log_var", "y_mask": "valid_mask"}}
    ),
    Regularizer(penalty="identity").set_io(
        {"inputs": {"x": "log_var", "y_mask": "valid_mask"}}
    ),
])

YAML:

loss:
  _target: srforge.loss.LossCombiner
  params:
    losses:
      - _target: srforge.loss.metrics.UncertaintyLoss
        params:
          base_metric:
            _target: srforge.loss.metrics.MSE
            params: {}
        io:
          inputs: {x: sr, y: hr, log_var: log_var, y_mask: valid_mask}
      - _target: srforge.loss.metrics.Regularizer
        params: {penalty: identity}
        io:
          inputs: {x: log_var, y_mask: valid_mask}

The base must implement pointwise() (Level 2). Trying to wrap a structural loss like SSIM or LPIPS fails loudly at construction.

Multi-task with shared uncertainty¶

The motivating use case for the split: applying uncertainty weighting to two or more base losses sharing a single σ field. Just add multiple UncertaintyLoss instances and exactly one Regularizer:

loss:
  _target: srforge.loss.LossCombiner
  params:
    losses:
      - _target: srforge.loss.metrics.UncertaintyLoss
        params: { base_metric: { _target: srforge.loss.metrics.MGE, params: {} } }
        io: { inputs: {x: sr, y: hr, log_var: log_var, y_mask: valid_mask} }
      - _target: srforge.loss.metrics.UncertaintyLoss
        params: { base_metric: { _target: srforge.loss.metrics.MSE, params: {} } }
        io: { inputs: {x: sr, y: hr, log_var: log_var, y_mask: valid_mask} }
      - _target: srforge.loss.metrics.Regularizer            # ← only ONCE
        params: { penalty: identity }
        io: { inputs: {x: log_var, y_mask: valid_mask} }

Math:

total = (e_MGE + e_MSE) · exp(-σ) + σ

Compare to the bug-prone alternative of stacking two old-style UncertaintyLoss instances (each carrying its own +σ term) — you'd get 2σ and the network would learn a systematically smaller uncertainty than it should.

Forgetting the Regularizer

If you use UncertaintyLoss without a companion Regularizer, σ has no penalty for growing and the network will push it toward +∞ (loss collapses). Manifests fast in training — the loss drops to large negative values within a few epochs. Always pair them.

Regularizer¶

Single-tensor magnitude penalty — generic regularizer. Reduces one Entry field to a per-sample scalar via a pointwise penalty function + the standard reduce(). Useful whenever an auxiliary output of your model needs a magnitude penalty.

Built-in penalties¶

`penalty=`	Pointwise function	Typical use
`"identity"`	`x`	Log-variance regularizer (companion to `UncertaintyLoss`)
`"abs"`	`\|x\|`	L1 / sparsity penalty (sparse attention, latent codes)
`"square"`	`x²`	L2 / magnitude penalty (predicted offsets, residuals)
`"exp"`	`exp(x)`	Penalise positive log-quantities
`"entropy"`	`-x·log(x)`	Element-wise entropy (reduce over softmax dim → categorical entropy)
`"huber"`	smooth-L1 with `huber_delta`	Quadratic near 0, linear far from it; robust magnitude
callable	`f(x)`	Anything custom — must return same-shape tensor

from srforge.loss.metrics import Regularizer

# L1 sparsity on attention weights
Regularizer(penalty="abs").set_io({"inputs": {"x": "attention_map"}})

# L2 magnitude penalty on a learned offset map
Regularizer(penalty="square").set_io({"inputs": {"x": "noise_pred", "y_mask": "valid_mask"}})

# Huber with custom delta
Regularizer(penalty="huber", huber_delta=0.5)

# Custom — log-L1 (robust soft-L1)
Regularizer(penalty=lambda x: torch.log1p(x.abs()))

YAML:

- _target: srforge.loss.metrics.Regularizer
  params: { penalty: abs }
  io:
    inputs: {x: attention_map}

best_min is always True — regularizers are always minimised.

LossScheduler¶

Switch between loss functions at epoch milestones:

from srforge.loss import L1, SSIM, LossCombiner
from srforge.loss.schedule import LossScheduler

scheduler = LossScheduler(schedule={
    0: LossCombiner([
        L1(weight=1.0).set_io({"inputs": {"x": "sr", "y": "hr"}}),
        SSIM(weight=0.01).set_io({"inputs": {"x": "sr", "y": "hr"}}),
    ]),
    50: LossCombiner([
        L1(weight=1.0).set_io({"inputs": {"x": "sr", "y": "hr"}}),
        SSIM(weight=0.05).set_io({"inputs": {"x": "sr", "y": "hr"}}),
    ]),
})

# During training:
scheduler.update(epoch)
scores = scheduler(entry)

The schedule must contain key 0. Call update(epoch) each epoch to switch the active loss.

YAML Configuration¶

In config files, use io: (same as transforms and models):

loss:
  _target: srforge.loss.schedule.LossScheduler
  params:
    schedule:
      0:
        _target: srforge.loss.combiner.LossCombiner
        params:
          losses:
            - _target: srforge.loss.metrics.L1
              params:
                weight: 1.0
              io:
                inputs: {x: sr, y: hr}
            - _target: srforge.loss.metrics.SSIM
              params:
                weight: 0.01
              io:
                inputs: {x: sr, y: hr}
      50:
        _target: srforge.loss.combiner.LossCombiner
        params:
          losses:
            - _target: srforge.loss.metrics.L1
              params:
                weight: 1.0
              io:
                inputs: {x: sr, y: hr}
            - _target: srforge.loss.metrics.SSIM
              params:
                weight: 0.05
              io:
                inputs: {x: sr, y: hr}

Next: Writing Scripts — Training and test scripts that wire everything together