SR-Forge

Extending SR-Forge

Every component in SR-Forge follows the same pattern: subclass, implement method(s), use from Python or YAML. This page shows how to create custom components — models, transforms, metrics (and the losses built from them), datasets, and hooks — that integrate seamlessly with the framework. Your custom classes work exactly like built-in ones: they can be wired in YAML, composed in SequentialModel, and used in any pipeline.

---

Model

Subclass Model, implement forward(). Parameter names are used for IO binding automatically.

from srforge.models import Model
import torch

class MyUpscaler(Model):
    def __init__(self, scale: int = 2):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Conv2d(3, 64, 3, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(64, 3 * scale ** 2, 3, padding=1),
            torch.nn.PixelShuffle(scale),
        )

    def forward(self, image):
        return self.net(image)

Use in YAML:

model:
  _target: my_project.models.MyUpscaler
  params:
    scale: 4
  io:
    inputs:
      image: lr
    outputs: sr

See Model for details on multiple outputs and standalone use.

---

DataTransform

Subclass DataTransform, implement transform() (batched) or transform_unbatched() (per-sample). Parameter names and type annotations define the interface.

from srforge.transform import DataTransform
import torch

class GammaCorrection(DataTransform):
    def __init__(self, gamma: float = 2.2):
        super().__init__()
        self.gamma = gamma

    def transform(self, image: torch.Tensor) -> torch.Tensor:
        return image.clamp(min=0).pow(1.0 / self.gamma)

Use in YAML:

preprocessing:
  training:
    - _target: my_project.transforms.GammaCorrection
      params:
        gamma: 2.2
      io:
        inputs:
          image: lr

Or in a SequentialModel flow:

modules:
  gamma:
    _target: my_project.transforms.GammaCorrection
    params:
      gamma: 2.2
flow:
  - "lr -> gamma -> lr"

See DataTransform for annotation-driven recursion, multi-input transforms, and per-sample processing.

---

EntryTransform

Subclass EntryTransform, implement transform_unbatched() (or transform() for batched). Use _key suffix parameters for field names.

from srforge.transform import EntryTransform

class DropSmallFields(EntryTransform):
    def __init__(self, *, min_size: int, field_key: str, output_key: str = None):
        self.min_size = min_size
        self.field_key = field_key
        self.output_key = output_key or field_key
        super().__init__()

    def transform_unbatched(self, entry):
        tensor = entry[self.field_key]
        if tensor.shape[-1] >= self.min_size and tensor.shape[-2] >= self.min_size:
            entry[self.output_key] = tensor
        return entry

Use in YAML — field keys go in params:, not io::

- _target: my_project.transforms.DropSmallFields
  params:
    min_size: 32
    field_key: image

See EntryTransform for the _key convention, optional keys, and SequentialModel integration.

---

Metric

Subclass Metric and provide the best_min property, then implement whichever of the three stages apply. A loss is a metric you train on, so this is also how you write a loss. The full guide is Writing a Metric: three questions that tell you which methods to write, and one metric written every possible way. This section is the short version. Most metrics need exactly one of the first two:

Level 2 — override pointwise() only (preferred for new metrics)

Return a per-element score map; the framework's default calculate_score composes reduce(pointwise(...), mask=y_mask) for you. The metric's inputs are pointwise's parameters — any number, any names — plus y_mask, which the framework adds. Masking, batch reduction, and per-sample normalisation are all handled automatically. Your metric also becomes wrappable by UncertaintyLoss(base_metric=...) and CorrectedLoss(base_metric=...).

from srforge.metrics import Metric
import torch

class HuberLoss(Metric):
    def __init__(self, delta: float = 1.0, **kwargs):
        super().__init__(**kwargs)          # forward framework kwargs!
        self.delta = delta

    @property
    def best_min(self) -> bool:
        return True

    def pointwise(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        diff = (x - y).abs()
        quadratic = torch.clamp(diff, max=self.delta)
        linear = diff - quadratic
        return 0.5 * quadratic ** 2 + self.delta * linear     # [B, C, H, W]

That's it. The framework reduces and masks automatically.

Level 1 — override calculate_score() (structural metrics)

Use this when the metric has no per-pixel value (SSIM uses windowed convolutions, MGE and TotalVariation compare neighbouring pixels), or when something happens after the average (PSNR takes a logarithm). Return a [B] tensor, and apply or refuse y_mask yourself.

class StructuralLoss(Metric):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @property
    def best_min(self) -> bool:
        return True

    def calculate_score(self, x: torch.Tensor, y: torch.Tensor,
                        y_mask: torch.Tensor = None) -> torch.Tensor:
        # Custom windowed/structural math; must return shape [B].
        # Every parameter needs a type annotation.
        ...
        return per_sample_scalar

For windowed metrics, compute the score map on the real inputs and average it over self.erode_mask(mask, window_size), which keeps only positions whose whole window is valid. Don't zero the masked inputs before the window runs: the zeros end up inside neighbouring windows (see Using Metrics → Masking).

The epoch stage — override finalize() (optional, either level)

Independent of the two above. By default the epoch value is the average of the per-image scores; override finalize when it isn't — a ratio, a logarithm, a total, a worst case. Declare what each part is so the framework combines it correctly over the epoch and across processes:

class MicroF1(Metric):
    @property
    def best_min(self) -> bool:
        return False

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

    def calculate_score(self, x: torch.Tensor, y: torch.Tensor) -> dict:
        ...
        return {"tp": tp, "fp": fp, "fn": fn}      # ingredients, per image

    def finalize(self, raw):
        return 2 * raw["tp"] / (2 * raw["tp"] + raw["fp"] + raw["fn"])

Valid reductions are sum, mean, min, max, cat. A Level-2 metric that overrides finalize receives {"sum", "count"} and needs no pixel arithmetic of its own — see Writing a Metric → Epoch rules.

Overriding finalize like this makes a metric custom: one fixed epoch rule, no aggregate option, not part of the training total.

Offering micro and macro — declare aggregations

A Level-2 metric gets aggregate="micro" for free. A Level-1 metric has to say how it pools — here is PSNR's version, slightly condensed:

class PSNR(Metric):
    aggregations = ("macro", "micro")          # first one is the default

    def calculate_score(self, x, y, y_mask=None):
        squared_error = torch.square(y - x)
        if self.aggregate == "micro":
            return self._sum_and_count(squared_error, y_mask)   # ingredients
        return self._psnr(self.reduce(squared_error, mask=y_mask))

    def finalize(self, raw):                   # only called under micro
        return self._psnr(raw["sum"] / raw["count"])

Three rules:

• Under micro, return {"sum", "count"} (from _sum_and_count) and the framework declares their reductions for you. Other ingredients need a reductions property, as in MicroF1 above.
• finalize runs over one batch during training and over the epoch for reporting, so write it for totals, not for a particular scale.
• Declaring "micro" without a way to produce it — a Level-1 metric with no finalize — fails when the class is defined, not when it is first used.

Set aggregation_note = "..." to explain, in the refusal message, why your class does not support something.

For any per-pixel yes/no metric — segmentation, change detection, anomaly detection — subclass ConfusionRatio rather than writing this by hand: it supplies the counts, the masking and both averagings. See Classification Metrics.

Framework kwargs — accept ALL of them, and forward

Metric.__init__ is decorated with @audit_subclasses, which enforces two rules at class definition (both raise TypeError):

1. Coverage — your subclass must ACCEPT every framework kwarg (weight, name, reduction, aggregate), either via **kwargs (recommended) or by declaring each explicitly. YAML users can't configure params your signature doesn't take.
2. Forwarding — whatever you accept must be passed to super. Either super().__init__(**kwargs) or super().__init__(weight=weight, name=name, reduction=reduction, aggregate=aggregate).

Simplest compliant pattern: def __init__(self, my_param, **kwargs): super().__init__(**kwargs).

Use in YAML:

loss:
  _target: srforge.metrics.LossCombiner
  params:
    losses:
      - _target: my_project.metrics.HuberLoss
        params:
          delta: 0.5
          weight: 1.0
        io:
          inputs: {x: sr, y: hr}

Need a custom regularizer? Use Regularizer, don't subclass

If your metric is a single-tensor magnitude penalty (penalise the mean / L1 / L2 / entropy of one Entry field), don't write a custom Metric subclass — use the built-in :class:Regularizer directly:

import torch
from srforge.metrics.regularization import Regularizer

# Custom robust soft-L1 penalty on attention weights
reg = Regularizer(penalty=lambda x: torch.log1p(x.abs()))
reg.set_io({"inputs": {"x": "attention_map"}})

Regularizer accepts any callable Tensor → Tensor and handles masking + reduction for you. Only subclass Metric for a regularizer if you need multi-tensor inputs or stateful logic. See Using Metrics → Regularizer for the built-in penalty options.

See Using Metrics for masking, multi-band inputs, MetricScores, and LossCombiner.

---

Dataset

Subclass Dataset from srforge.dataset and implement __getitem__ and __len__. Each __getitem__ call returns one Entry. The framework auto-wraps your __getitem__ (via Dataset.__init_subclass__) to apply the constructor transforms= list and to handle disk caching when cache_dir= is set.

Framework kwargs — accept ALL of them

Same rule as for Metric: Dataset is decorated with @audit_subclasses, so your subclass must accept every framework kwarg (name, transforms, cache_dir, recache) — either via **kwargs or by declaring each explicitly. Forgetting raises TypeError at class definition. Simplest compliant pattern:

def __init__(self, my_param, **kwargs):
    super().__init__(**kwargs)

from srforge.dataset import Dataset
from srforge.data import Entry
from pathlib import Path
import torch

class TiffDataset(Dataset):
    def __init__(self, root: str, **kwargs):
        super().__init__(**kwargs)  # accepts name=, transforms=, cache_dir=, recache=
        self._paths = sorted(Path(root).glob("*.tif"))

    def __getitem__(self, index: int) -> Entry:
        path = self._paths[index]
        image = load_tiff(path)  # your loading function
        return Entry(
            name=path.stem,
            image=torch.from_numpy(image).float(),
        )

    def __len__(self) -> int:
        return len(self._paths)

Use in YAML:

dataset:
  training:
    _target: my_project.datasets.TiffDataset
    params:
      root: /data/train
      cache_dir: /tmp/tiff-cache   # optional — caches transformed entries
      transforms:                   # applied automatically per __getitem__
        - _target: srforge.transform.data.ZScore
          io: {inputs: {x: image}, outputs: image}

Datasets support transforms at load time and on-disk caching — see Datasets. Useful operators (take, filter, shuffle, + for concat) are inherited from the base class.

---

Hook

Subclass Hook and decorate handler methods with @hooks_into("on_<point>") — the decorator says which HookPoint on the runner/trainer the method fires at:

from srforge.training.hooks import Hook, hooks_into

class EpochTimer(Hook):
    @hooks_into("on_epoch_end")
    def print_time(self, ctx):
        import time
        print(f"Epoch {ctx.epoch} finished at {time.strftime('%H:%M:%S')}")

Use in YAML — put it in the hooks: list of the component whose lifecycle it follows (here: the training runner):

training_runner:
  _target: srforge.training.runners.TrainingEpochRunner
  params:
    # ... optimizer, postprocessor, ...
    hooks:
      - _target: my_project.hooks.EpochTimer

The ctx argument is a mutable Context — hooks can read entry data, scores, and epoch state, and can also modify training (add auxiliary losses, clip gradients, skip steps). See Hooks for the full HookPoint list, the Context API, and stateful hooks.

Custom observers were removed in 0.16.0

Code written for 0.15.x and earlier may subclass Observer with an EVENTS list. srforge.observers no longer exists; turn each one into a hook with the upgrading guide — the conversion is mechanical.

---

Making Classes Available in YAML

All custom classes are used in YAML via their full module path:

_target: my_project.models.MyUpscaler

The ConfigResolver imports the module and instantiates the class automatically — no registration step needed. As long as the module is importable (i.e., on sys.path or installed as a package), it works.

For a typical project layout:

my-experiment/
├── train.py
├── configs/
│   └── train-cfg.yaml
└── my_project/
    ├── __init__.py
    ├── models.py          # _target: my_project.models.MyUpscaler
    ├── transforms.py      # _target: my_project.transforms.GammaCorrection
    ├── metrics.py          # _target: my_project.metrics.HuberLoss
    ├── datasets.py         # _target: my_project.datasets.TiffDataset
    └── hooks.py            # _target: my_project.hooks.EpochTimer

SR-Forge built-in classes use the srforge.* prefix:

_target: srforge.metrics.regression.L1
_target: srforge.transform.data.Multiply
_target: srforge.training.trainers.PyTorchTrainer

---

Next: Return to Getting Started to scaffold a project, or browse the API Reference for detailed class documentation.