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, losses, 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.

Loss¶

Subclass Loss and provide the best_min property. You have two implementation styles:

Level 2 — override `pointwise()` only (preferred for new losses)¶

Return a per-element score map; the framework's default calculate_score composes reduce(pointwise(x, y), mask=y_mask) for you. Masking, batch reduction, and per-sample normalisation are all handled automatically. Your loss also becomes wrappable by UncertaintyLoss(base_metric=...) and CorrectedLoss(base_metric=...).

from srforge.loss import Loss
import torch

class HuberLoss(Loss):
    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 losses)¶

Use this when the loss isn't decomposable into a per-element op — e.g. SSIM uses windowed convolutions, LPIPS uses a learned feature network, SAM computes spectral angles. Return a [B] tensor.

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

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

    def calculate_score(self, x, y, y_mask=None):
        # Custom windowed/structural math; must return shape [B].
        ...
        return per_sample_scalar

For windowed losses that need masked-position neutralisation before the windowed op, use the mask_pixels() helper (see Losses → Masking).

Framework kwargs — accept ALL of them, and forward

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

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

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

Use in YAML:

loss:
  _target: srforge.loss.LossCombiner
  params:
    losses:
      - _target: my_project.losses.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 loss is a single-tensor magnitude penalty (penalise the mean / L1 / L2 / entropy of one Entry field), don't write a custom Loss subclass — use the built-in :class:Regularizer directly:

import torch
from srforge.loss.metrics 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 Loss for a regularizer if you need multi-tensor inputs or stateful logic. See Losses → Regularizer for the built-in penalty options.

See Losses 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 Loss: 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 are deprecated

Older code subclassed Observer with an EVENTS list and subscribed to the global EventBus. That mechanism is deprecated — removal in v0.16.0 and emits a DeprecationWarning at both class definition and instantiation. Port existing observers with the Observer → Hook 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
    ├── losses.py           # _target: my_project.losses.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.loss.metrics.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.

Extending SR-Forge¶

Model¶

DataTransform¶

EntryTransform¶

Loss¶

Level 2 — override pointwise() only (preferred for new losses)¶

Level 1 — override calculate_score() (structural losses)¶

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

Dataset¶

Hook¶

Making Classes Available in YAML¶

Level 2 — override `pointwise()` only (preferred for new losses)¶

Level 1 — override `calculate_score()` (structural losses)¶

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