SR-Forge

Model¶

A Model is a PyTorch neural network that participates in SR-Forge's pipeline. You write forward() with your computation — the framework handles extracting inputs from Entry fields and storing outputs back. It works exactly like a regular nn.Module, with IO binding on top.

Your First Model¶

from srforge.models import Model
import torch

class Upscaler(Model):
    def __init__(self, channels: int = 64):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Conv2d(3, channels, 3, padding=1),
            torch.nn.ReLU(),
            torch.nn.Conv2d(channels, 3, 3, padding=1),
        )

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

Three things to notice:

Inputs are automatic — inferred from forward() parameter names (here: image). Outputs are determined when you bind the model (via set_io() or the flow DSL).
forward() — your computation, just like any PyTorch module. Parameter names are used for IO binding automatically.
It's a regular nn.Module — you can call it standalone with plain tensors: model(some_tensor)

The `forward()` Convention¶

In standard PyTorch, you override forward() in your subclass. SR-Forge follows the same convention:

class MyModel(Model):
    def forward(self, x):
        return self.net(x)

This is all you need to know: just define forward() as in any PyTorch module — the framework handles the rest invisibly.

Behind the scenes: the transparent rename

Model.__init_subclass__ renames your forward() to _forward() at class-definition time, so the base class's own Model.forward (the Entry-routing layer) can intercept calls, extract fields per your IO binding, call your code, and store outputs back. This keeps the standard PyTorch forward() authoring convention while still giving the framework its routing layer.

The older _forward() authoring convention still works for backward compatibility:

class MyModel(Model):
    def _forward(self, x):
        return self.net(x)

Both produce identical behavior. Defining both forward and _forward raises TypeError.

Inputs and Outputs¶

Model inputs are always inferred from the forward() signature — no declaration needed:

Parameters without defaults become required inputs
Parameters with defaults become optional inputs — they can be omitted from the IO config, and the default value applies

class MyModel(Model):
    def forward(self, image, guide=None):  # image: required, guide: optional
        ...

If a parameter is not mapped in the IO config (omitted from set_io() or flow DSL), its default value is used. But if a parameter is mapped to an Entry field, that field must exist — a KeyError is raised otherwise, even if the parameter has a default. This catches misconfigured pipelines where a field you expected is missing.

Output field names come from one of two sources:

set_io() — the outputs key specifies where results are written
Flow DSL — positional mapping assigns output fields directly

You never need to declare inputs or outputs ahead of time — inputs are inferred from forward() parameters, and outputs are specified at binding time.

How Entry Routing Works¶

When you call set_io(), the inputs dict maps forward() parameter names to Entry fields, and outputs specifies where results are written. When you call model(entry), the base class routes data through the binding:

Detects that the argument is an Entry
Looks up the IO binding: image -> "image_rgb"
Extracts entry["image_rgb"] and passes it as image=entry["image_rgb"]
Your forward() runs and returns a tensor
Maps the return value to the output field: "prediction"
Writes entry["prediction"] = result
Returns the updated Entry

model = Upscaler()
model.set_io({"inputs": {"image": "image_rgb"}, "outputs": "prediction"})
# "image" → forward() param, reads entry["image_rgb"]
# "prediction" → output, writes entry["prediction"]

entry = Entry(image_rgb=some_tensor)
result = model(entry)
# result.prediction contains the model output
# result.image_rgb is preserved

Multiple Outputs¶

If your model returns multiple values, map each to a field using a list or positional flow DSL:

With set_io() — list of output field names:

class FusionModel(Model):
    def forward(self, image, guide):
        output, attn = self.fusion_net(image, guide)
        return output, attn

model = FusionModel()
model.set_io({
    "inputs": {"image": "x", "guide": "edges"},
    "outputs": ["result", "attn"],
})
# 1st return value → entry["result"]
# 2nd return value → entry["attn"]

With flow DSL — positional mapping:

flow = ["(x, edges) -> fuse -> (result, conf)"]
# 1st return value -> entry["result"]
# 2nd return value -> entry["conf"]

Field Overwrite Protection¶

Model outputs cannot overwrite existing Entry fields — this prevents accidentally losing data (e.g., a model silently replacing the ground truth). If entry["prediction"] already exists, a KeyError is raised:

entry = Entry(image_rgb=tensor, prediction=existing_tensor)
model(entry)  # KeyError: attempted to overwrite existing entry fields: ['prediction']

Standalone Use¶

Models can be called directly with raw tensors, bypassing IO binding entirely:

# Pipeline mode (with Entry)
model.set_io({"inputs": {"image": "image_rgb"}, "outputs": "prediction"})
result_entry = model(entry)

# Standalone mode (plain PyTorch)
output = model(image=some_tensor)  # Returns raw forward() output

In standalone mode, the return value is whatever forward() returns — no Entry wrapping, no output mapping. No set_io() needed.

IO Binding¶

Model IO is configured through set_io(), YAML io: keys, or the SequentialModel flow DSL.

In Python¶

model.set_io({"inputs": {"image": "image_rgb"}, "outputs": "prediction"})

In YAML Config¶

model:
  _target: Upscaler
  params:
    channels: 128
  io:
    inputs:
      image: image_rgb
    outputs: prediction

In SequentialModel¶

seq = SequentialModel(
    modules={"sr": Upscaler()},
    flow=["x -> sr -> y"],
)
# Parameter "image" reads entry["x"], output writes entry["y"]

For the full catalog of binding syntax, see the IO Binding Reference: Model.

In SequentialModel¶

Models are the most flexible module type inside a SequentialModel. The flow DSL controls which Entry fields map to which parameters:

flow = ["(x, edges) -> enhance -> (result, conf)"]

For explicit input mapping, use named syntax on the LHS:

flow = ["(image=x, guide=edges) -> enhance -> (result, conf)"]

Models are fully reusable — the same instance can appear in multiple flow lines with shared weights:

encoder = Encoder()
seq = SequentialModel(
    modules={"encoder": encoder},
    flow="""
        x -> encoder -> h1
        h1 -> encoder -> h2
    """,
)
# encoder is applied twice: x -> h1 -> h2 (shared weights)

See SequentialModel: Models in Flow for details.

Loading pretrained weights¶

The canonical config-callable factory for restoring W&B-tracked model weights is srforge.utils.checkpoint.load_weights_from_wandb. Use it as a _target wrapping your bare model — the resolver instantiates the model first, then this function loads the state-dict in place and returns the same instance:

model:
  _target: srforge.utils.checkpoint.load_weights_from_wandb
  params:
    module:                                # the bare model
      _target: srforge.models.MISR.RAMS.RAMS
      params: { scale: 3 }
    project: my-project
    entity:  my-entity
    run_id:  abc123xyz
    load_best_model: true                  # picks {ModuleName}_best.pth
    # filename / state_key / strict / legacy_weight_norm — see srforge/utils/checkpoint.py
  io:
    inputs:  { lrs: lrs }
    outputs: sr

For full parameters, alternative-run loading, and the in-script variant load_weights_from_tracker, see Experiment Tracking § Loading weights from a different run.

Next: GANModel — Adversarial training with generator and discriminator | SequentialModel — Compose models and transforms into multi-stage pipelines

Model¶

Your First Model¶

The forward() Convention¶

Inputs and Outputs¶

How Entry Routing Works¶

Multiple Outputs¶

Field Overwrite Protection¶

Standalone Use¶

IO Binding¶

In Python¶

In YAML Config¶

In SequentialModel¶

In SequentialModel¶

Loading pretrained weights¶

The `forward()` Convention¶