SR-Forge

Architecture

Visual overview of SR-Forge's structure and how components connect.

---

Core Concepts

The pieces that matter at runtime and how they fit together. Solid arrows are direct calls or ownership; dashed arrows are the config hand-off (cfg) and hook firing. ConfigResolver instantiates everything from a YAML config; Trainer owns the lifecycle; EpochRunner executes one epoch at a time over the pipeline, while Hooks (attached directly to the trainer and runners) record — and can influence — what happens, logging through the Tracker.

flowchart TB
    INIT(["srforge.init(cfg)"]):::entry
    YAML[/"experiment.yaml<br/><small>_target, params, io, …</small>"/]:::yaml
    INIT ~~~ YAML
    YAML -. cfg .-> INIT
    CONFIG["ConfigResolver"]:::accent
    INIT -->|builds| CONFIG
    CONFIG -->|instantiates| TRAINER["Trainer"]
    TRAINER -->|owns| RUNNER["EpochRunner<br/><small>training · validation</small>"]

    subgraph pipeline["Per-batch pipeline"]
        direction LR
        DATASET["Dataset"]
        MODEL["Model"]
        LOSS["Loss"]
    end

    RUNNER -->|drives| pipeline

    subgraph side["Side channel"]
        direction TB
        HOOKS["Hooks<br/><small>ProgressBar · LossLogger ·<br/>ModelSaver · GradientClip · …</small>"]
        TRACKER["Tracker"]
        HOOKS -->|log metrics / images| TRACKER
    end

    TRAINER -. "fires HookPoints" .-> HOOKS
    RUNNER -. "fires HookPoints" .-> HOOKS
    TRAINER -->|saves / restores| TRACKER
    RUNNER -->|saves state| TRACKER

    subgraph foundation["Foundation"]
        direction LR
        ENTRY["Entry"]:::accent
        IOM["IOModule"]:::accent
    end

    pipeline -->|reads / writes Entry| foundation

    classDef accent stroke:#c9a84c,stroke-width:2.5px;
    classDef entry stroke:#c9a84c,stroke-width:2.5px,stroke-dasharray:4 3;
    classDef yaml stroke:#6b7280,stroke-width:1.5px,stroke-dasharray:2 2;

Where did the EventBus go?

Older versions routed monitoring through a global EventBus + Observer system. That mechanism is deprecated (removal in v0.16.0) and survives only via a compatibility adapter — see Observers & Events. New code attaches Hooks directly to the trainer / runners.

A few things this diagram intentionally doesn't show — they live in their own pages:

• Per-class inheritance (e.g. Model(IOModule, nn.Module, ABC)) — see the Class Hierarchy diagram below.
• Detailed step-by-step batch flow — see the Training Flow sequence diagram below.

---

Class Hierarchy

Inheritance relationships for the major class families. SR-Forge mixes three orthogonal capabilities into its base classes:

• IOModule (srforge.utils.iomodule.IOModule) — adds the set_io() / io: field-binding machinery.
• torch.nn.Module — parameter / state-dict / .train() / .eval() semantics. Required because Model and Loss participate in autograd.
• Observable (srforge.observers.Observable) — deprecated event-bus emitter. Trainer / EpochRunner no longer inherit it (they expose :class:HookPoint attributes instead and register with the legacy bus directly). User code that still subclasses Observable keeps working with a DeprecationWarning.

Model(IOModule, torch.nn.Module, abc.ABC) and Loss(nn.Module, IOModule, ABC) are the two canonical multiple-inheritance chains carried into the autograd graph. Trainer(abc.ABC) and EpochRunner(abc.ABC) no longer mix in Observable — they expose lifecycle stages as :class:HookPoint attributes (on_training_began, on_post_step, …) that hooks register with directly. The data pipeline shares the same IOModule mixin: DataTransform(IOModule, Transformation, abc.ABC) and EntryTransform(IOModule, Transformation, abc.ABC) bind to Entry fields the way models and losses do. The diagram covers six families:

• Data pipeline — Dataset (yields Entry objects), Transformation → DataTransform / EntryTransform, and the Entry primitive itself (_DynamicStorage subclass).
• Models — Model → GANModel / SequentialModel and concrete architectures.
• Training — Trainer → PyTorchTrainer, EpochRunner → the training / validation / benchmark / GAN runners.
• Side-channel: hooks — HookPoint + Handle are first-class stage points; Hook subclasses (StepRatio, GradientClip, LossLogger, …) plug into them. The legacy Observer family is shown as deprecated.
• Losses & metrics — Loss → LossCombiner and metric losses (L1, MSE, SSIM, …), producing MetricScores (which hold MetricEntry rows).
• Tracking & observers — ExperimentTracker and Observer families.

%%{init: {"class": {"hideEmptyMembersBox": true}}}%%
classDiagram
    direction TB

    %% Pipeline mixins — kept adjacent to their consumers (Model, Loss,
    %% Transforms) so the inheritance edges stay short.
    class IOModule {
        <<mixin>>
        +set_io(io_cfg)
        +io: IO
    }

    class nn_Module["torch.nn.Module"]

    class Model {
        <<abstract>>
        +forward(data) Any
        +trainable_params() List
        #_forward(*args, **kwargs)*
    }

    class GANModel {
        +generator: Model
        +discriminator: Model
        +discriminator_step(entry) Entry
        +generator_step(entry) Entry
    }

    class SequentialModel {
        +steps: List~Step~
        +models: ModuleDict
    }

    class moreModels["⋯ more"]:::more

    IOModule <|-- Model
    nn_Module <|-- Model
    Model <|-- GANModel
    Model <|-- SequentialModel
    Model <|-- moreModels

    class Loss {
        <<abstract>>
        +calculate_score(*args)*
        +weight: float
        +best_min: bool*
    }

    class LossCombiner {
        +losses: List~Loss~
    }

    class L1
    class MSE
    class SSIM
    class moreLoss["⋯ more"]:::more

    nn_Module <|-- Loss
    IOModule <|-- Loss
    Loss <|-- LossCombiner
    Loss <|-- L1
    Loss <|-- MSE
    Loss <|-- SSIM
    Loss <|-- moreLoss

    class MetricScores {
        +total_weighted() Tensor
        +mean() MetricScores
        +merge(other)
    }

    class MetricEntry {
        +value: Tensor
        +weight: float
        +best_min: bool
    }

    Loss ..> MetricScores : produces
    MetricScores o-- MetricEntry : holds

    class Transformation {
        <<abstract>>
        +__call__(data)
        +transform(data)*
    }

    class DataTransform {
        <<abstract>>
        +transform(**fields)
    }

    class EntryTransform {
        <<abstract>>
        +transform(entry) Entry
    }

    class ZScore
    class Downsample
    class Normalize
    class RandomCrop
    class StackBands
    class moreDT["⋯ more"]:::more
    class moreET["⋯ more"]:::more

    Transformation <|-- DataTransform
    Transformation <|-- EntryTransform
    IOModule <|-- DataTransform
    IOModule <|-- EntryTransform
    DataTransform <|-- ZScore
    DataTransform <|-- Downsample
    DataTransform <|-- Normalize
    DataTransform <|-- moreDT
    EntryTransform <|-- RandomCrop
    EntryTransform <|-- StackBands
    EntryTransform <|-- moreET

    %% Trainer / EpochRunner — no longer inherit Observable.
    class Trainer {
        <<abstract>>
        +train(epochs, train_loader, val_loader)
        +training_runner: EpochRunner
        +hooks: list~Hook~
    }

    class PyTorchTrainer

    Trainer <|-- PyTorchTrainer

    class EpochRunner {
        <<abstract>>
        +run_epoch(model, loader, epoch) MetricScores
        +training_state() dict
        +load_training_state(state)
        +zero_grad()
        +hooks: list~Hook~
    }

    class TrainingEpochRunner {
        +optimizer: Optimizer
        +scaler: GradScaler
        +gradient_accumulation_steps: int
    }

    class GANTrainingRunner {
        +optimizer_G: Optimizer
        +optimizer_D: Optimizer
        +d_criterion: Loss
        +g_criterion: Loss
    }

    class InferenceRunner["InferenceRunner · no backward"]
    class ValidationEpochRunner
    class BenchmarkRunner

    EpochRunner <|-- TrainingEpochRunner
    EpochRunner <|-- GANTrainingRunner
    EpochRunner <|-- InferenceRunner
    InferenceRunner <|-- ValidationEpochRunner
    InferenceRunner <|-- BenchmarkRunner

    %% Side-channel system: HookPoint + Hook (canonical).
    class HookPoint {
        +name: str
        +register(fn) Handle
        +fire(ctx)
    }

    class Handle {
        +remove()
    }

    HookPoint *-- Handle : creates

    class Hook {
        +scope: str | list | None
        +bind_to(target) list~Handle~
        +state_dict() dict
        +load_state_dict(state)
    }

    class StepRatio
    class R1GradientPenalty
    class GradientClip
    class moreHooks["⋯ more"]:::more

    Hook <|-- StepRatio
    Hook <|-- R1GradientPenalty
    Hook <|-- GradientClip
    Hook <|-- moreHooks

    Trainer ..> HookPoint : exposes
    EpochRunner ..> HookPoint : exposes
    Hook ..> HookPoint : binds_to

    %% Independent families — no cross-cutting mixin edges.
    class Dataset {
        <<torch.utils.data.Dataset>>
        +__getitem__(idx) Entry
        +__len__() int
    }

    class LazyDataset
    class PatchedDataset
    class LazyMultiSpectralDataset
    class moreDS["⋯ more"]:::more

    Dataset <|-- LazyDataset
    Dataset <|-- PatchedDataset
    Dataset <|-- moreDS
    LazyDataset <|-- LazyMultiSpectralDataset

    class _DynamicStorage["_DynamicStorage · abstract"]

    class Entry {
        +fields: dict
    }

    class GraphEntry["GraphEntry · PyG"]

    _DynamicStorage <|-- Entry

    class ExperimentTracker {
        <<abstract>>
        +log_metrics(dict, step)
        +save_file(path)
        +restore_file(filename)
    }

    class WandbTracker
    class NullTracker

    ExperimentTracker <|-- WandbTracker
    ExperimentTracker <|-- NullTracker

    %% Legacy Observer family — deprecated; ObserverAdapter wraps each
    %% subscribed Observer as a Hook so they keep working unchanged.
    class Observer {
        <<deprecated>>
        +EVENTS: list
    }

    class Observable {
        <<deprecated>>
    }

    class PyTorchModelSaver
    class LossLogger_legacy["LossLogger (legacy Observer)"]
    class ProgressBar_legacy["ProgressBar (legacy Observer)"]
    class moreObs["⋯ more"]:::more

    Observer <|-- PyTorchModelSaver
    Observer <|-- LossLogger_legacy
    Observer <|-- ProgressBar_legacy
    Observer <|-- moreObs

    class ObserverAdapter
    Hook <|-- ObserverAdapter
    ObserverAdapter ..> Observer : wraps

    classDef more fill:transparent,stroke:#6b7280,stroke-width:1px,stroke-dasharray:3 3,color:#9ca3af;

---

Training Flow

How data flows through the system during one training epoch. Solid arrows are calls; dashed arrows are returns. HookPoint firings (dotted notes) are where attached hooks — progress bars, loggers, checkpoint savers — get their turn.

sequenceDiagram
    participant T as PyTorchTrainer
    participant R as TrainingEpochRunner
    participant M as Model
    participant PP as Postprocessor
    participant L as Loss
    participant O as Optimizer
    participant H as Hooks

    T->>H: on_training_began
    Note over H: ProgressBar initialises,<br/>ModelSaver arms

    loop Each Epoch
        T->>R: run_epoch(model, loader, criterion)
        R->>H: on_epoch_start

        loop Each Batch
            R->>H: on_pre_step
            R->>M: forward(entry)
            M-->>R: entry with outputs (sr, etc.)
            R->>PP: postprocess(entry)
            PP-->>R: entry (maybe with cropped fields)
            R->>L: criterion(entry)
            L-->>R: MetricScores
            R->>R: loss.backward()
            R->>H: on_post_backward
            Note over H: GradientClip may act here
            R->>O: optimizer.step()
            R->>H: on_post_step
            Note over H: ProgressBar updates
        end

        R->>H: on_epoch_end
        R-->>T: epoch MetricScores

        T->>T: validation_runner.run_epoch(...)
        T->>H: on_trainer_epoch_finished
        Note over H: ModelSaver checkpoints,<br/>LossLogger logs metrics
        T->>T: lr_scheduler.step()
        T->>T: stop_condition.check()
    end

GAN Training Flow

The GAN runner adds hook dispatch and alternating D/G steps. Hook stages map directly onto HookPoint attributes on GANTrainingRunner — see Trainers & Runners → Training Hooks for the full HookPoint list and the Hook API.

sequenceDiagram
    participant R as GANTrainingRunner
    participant H as Hooks
    participant M as GANModel
    participant DL as D Loss
    participant GL as G Loss
    participant PL as Pixel Loss

    R->>H: dispatch(EPOCH_START)

    loop Each Batch
        R->>M: forward(entry) → generator produces SR

        R->>H: dispatch(PRE_STEP)
        Note over H: StepRatio sets run_d/run_g<br/>DWarmup may suppress G

        alt D Step (if ctx.run_d)
            R->>H: dispatch(PRE_D_FORWARD)
            Note over H: EdgeEnhancedInput<br/>modifies entry fields
            R->>M: discriminator_step(entry)
            M-->>R: entry with real_score_d, fake_score_d
            R->>DL: d_criterion(entry)
            DL-->>R: d_scores
            R->>H: dispatch(POST_D_FORWARD)
            Note over H: R1GradientPenalty<br/>adds to ctx.d_extra_losses
            R->>R: d_loss + extra_losses → backward
            R->>H: dispatch(POST_D_BACKWARD)
            Note over H: GradientClip on D
            R->>R: optimizer_D.step()
        end

        alt G Step (if ctx.run_g)
            R->>H: dispatch(PRE_G_FORWARD)
            R->>M: generator_step(entry)
            M-->>R: entry with real_score_g, fake_score_g
            R->>PL: pixel_criterion(entry)
            R->>GL: g_criterion(entry)
            R->>H: dispatch(POST_G_FORWARD)
            R->>R: pixel_loss + adv_loss + extra → backward
            R->>H: dispatch(POST_G_BACKWARD)
            R->>R: optimizer_G.step()
        end

        R->>H: dispatch(POST_STEP)
    end

    R->>H: dispatch(EPOCH_END)