SR-Forge

Experiment Tracking¶

SR-Forge provides an ExperimentTracker abstraction so your scripts and hooks are tracking-backend-agnostic. The backend (Weights & Biases, MLflow, or nothing) is selected via config.

The Interface¶

Every tracker implements ExperimentTracker:

Method	Description
`log_metrics(metrics, *, step)`	Log scalar metrics (`step` keyword-only, required)
`log_image(key, image, *, step)`	Log an image (`step` keyword-only, required)
`log_table(key, columns, data, *, step)`	Log a tabular artifact (default impl is a no-op; overridden by `WandbTracker`)
`set_summary(key, value)`	Set a run-level summary value
`commit(*, step)`	Flush pending logs at this step
`log_config(config)`	Log the experiment config
`watch_model(model, **kwargs)`	Attach gradient/parameter tracking
`save_file(path, *, base_path)`	Upload a file artifact (`base_path` keyword-only)
`restore_file(filename)`	Download a file from the current (resumed) run
`restore_from_run(run_id, filename)`	Download a file from a different run
`finish(exit_code)`	Finalize the run

Properties:

Property	Type	Description
`run_id`	`str \\| None`	Unique run identifier
`run_name`	`str \\| None`	Human-readable run name
`run_path`	`str \\| None`	Full path/URL to the run
`is_resumed`	`bool`	Whether this run was resumed

Built-in Trackers¶

WandbTracker¶

Backed by Weights & Biases. All constructor params are keyword-only and pass through to wandb.init() — extra kwargs forward via **kwargs, so any future wandb feature works without a docs update.

tracker:
  _target: srforge.tracking.WandbTracker
  params:
    project: my-project       # required
    entity: my-team           # optional — defaults to wandb's account
    mode: online              # online | offline | disabled
    name:                     # auto-generated if null
    run_id:                   # set to resume a previous run
    group:                    # group related runs
    tags: [exp-42, gan]       # list of tags
    notes: "GAN run …"        # free-text description
    job_type: training        # used by wandb's UI grouping
    resume: allow              # how wandb should resume — see below
    dir: ./wandb-runs          # local cache dir (default: ~/.cache/wandb)
    settings:                  # advanced — wandb.Settings instance
    save_code: true            # capture script + git info

Constructor arg	Default	Purpose
`project`	required	W&B project name
`entity`	`None`	Account / org; falls back to the wandb-CLI default
`name`	`None`	Display name for the run (auto-generated if `None`)
`run_id`	`None`	Resume the run with this id when it already exists
`group` / `tags` / `notes`	`None`	UI metadata
`job_type`	`"training"`	UI grouping for related run types
`resume`	`"allow"`	wandb resume mode (`"allow"` / `"must"` / `"never"`) — mutually exclusive with `fork_from` / `resume_from`
`mode`	`"online"`	`"online"` / `"offline"` / `"disabled"`
`dir`	`None`	Where wandb stores local run data
`settings`	`None`	A `wandb.Settings` instance for advanced tweaks
`save_code`	`True`	Snapshot the script + git state on `wandb.init`
`**kwargs`	—	Forwarded to `wandb.init` — use `fork_from=` or `resume_from=` here to branch from another run (these auto-disable `resume`)

Side effects on construction:

Calls wandb.init(**all_kwargs) — the network handshake happens here, not at first log.
Assigns GlobalSettings().tracker = self so the SIGINT handler in srforge.init() can call tracker.finish(1) on Ctrl+C.
Prints a Rich-styled run-info panel (Run ID / Name / Project / Path / Resumed: yes if applicable).

Resuming a run — when run_id is set and the run already exists, W&B resumes it. The is_resumed property returns True, which triggers checkpoint restoration in your script (see Resuming Training below).

Forking a run — use fork_from="<run_id>?_step=NNN" (a wandb spec; passed through **kwargs). The tracker drops the resume argument automatically when fork_from or resume_from is set, because wandb's pydantic validator considers them mutually exclusive.

log_table(key, columns, data, *, step) — in addition to the base API, WandbTracker overrides log_table to emit a wandb.Table for the run's UI. The base ExperimentTracker.log_table no-ops by default, so other backends ignore the call gracefully.

tracker.log_table(
    "predictions",
    columns=["sample", "image", "psnr"],
    data=[["s0", img_to_wandb(sr), 32.4], ...],
    step=epoch,
)

NullTracker¶

A no-op tracker where every method silently does nothing. Properties return None/False. Use it for offline development, debugging, or testing:

tracker:
  _target: srforge.tracking.NullTracker

restore_file() and restore_from_run() raise FileNotFoundError — there are no files to restore when there's no tracking backend.

NullTracker is also the default for hooks when no tracker is explicitly provided.

Usage in Scripts¶

The standard pattern:

from srforge import init
import omegaconf

resolve = init(cfg)

tracker = resolve(cfg.tracker)
tracker.log_config(omegaconf.OmegaConf.to_container(cfg, resolve=False))

# ... training ...

tracker.finish(0)

Hooks receive the tracker as a constructor parameter — they don't need to know which backend is active. In YAML, use ${ref:tracker} to inject the already-resolved tracker:

trainer:
  _target: srforge.training.trainers.PyTorchTrainer
  params:
    # ... runners, criteria, scheduler ...
    hooks:
      - _target: srforge.training.hooks.LossLogger
        params:
          tracker: ${ref:tracker}   # injects the resolved tracker instance
      - _target: srforge.training.hooks.PyTorchModelSaver
        params:
          tracker: ${ref:tracker}

Inside a hook, access is via self.tracker:

# Inside a hook — works with any tracker
self.tracker.log_metrics({"train/loss": loss_value}, step=epoch)
self.tracker.save_file(checkpoint_path, base_path=output_dir)

Resuming Training¶

Resume is driven by the tracker. When tracker.is_resumed is True, the resume_from_checkpoint() utility:

Downloads the last checkpoint via tracker.restore_file("checkpoint_last.pth")
Restores the runner's training state (optimizer + AMP scaler) via runner.load_training_state()
Restores LR-scheduler state via lr_scheduler.load_state_dict()
Loads model weights for every sub-model
Restores RNG states for exact reproducibility

from srforge.utils.checkpoint import resume_from_checkpoint

# Create the trainer FIRST — its `training_runner` exposes the
# optimizer + AMP scaler state that resume_from_checkpoint restores.
trainer = resolve(cfg.trainer)

ckpt = resume_from_checkpoint(
    model, trainer.training_runner, scheduler, tracker=tracker
)
trainer.restore(ckpt)  # None-safe: no-op for fresh runs

Signature changed

The 2^nd positional argument is the runner, not the optimizer. Older sr-forge code passed optimizer here; the function now calls runner.load_training_state(state) so it needs the runner the trainer created (trainer.training_runner).

For W&B, set run_id in the config to the run you want to resume:

tracker:
  _target: srforge.tracking.WandbTracker
  params:
    project: my-project
    run_id: abc123xyz      # this triggers resume
    resume: allow

Loading weights from a different run¶

Two helpers, picked by when you have the tracker and which W&B parameters you have to supply:

`load_weights_from_wandb` — the YAML-friendly factory¶

This is the canonical config-callable factory: it builds the W&B API client itself from project + entity + run_id, downloads the state-dict, loads it, and returns the populated module. Because it takes the module as its first positional argument and returns the same module, it works as a _target inside an io:-binding config node:

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 {module_name}_best.pth
    # filename: model_v2.pth                    # optional explicit override
    # state_key: model                          # if the .pth wraps the state-dict
    # strict:    false                          # optional non-strict load
    # legacy_weight_norm: true                  # opt-in compat for old checkpoints
  io:
    inputs:  { lrs: lrs }
    outputs: sr

The resolver instantiates module first, then calls load_weights_from_wandb(module, ...). The returned (now-populated) module is what model: resolves to — io: bindings on the outer node attach to it. See srforge/utils/checkpoint.py:274 for the full signature.

`load_weights_from_tracker` — for scripts with an existing tracker¶

When you're already inside a script that has a live WandbTracker and you want to overlay weights from a different run mid-flight:

from srforge.utils.checkpoint import load_weights_from_tracker

model = load_weights_from_tracker(
    module=model,
    tracker=tracker,
    run_id="other-run-id",
    best=True,
    module_name=None,   # defaults to module.__class__.__name__
)

Downloads the weight file via tracker.restore_from_run(...), loads it into the module, cleans up the temp file. Prefer load_weights_from_wandb for anything driven from YAML; reach for this one only when you already have a tracker object in hand.

Writing a Custom Tracker¶

Subclass ExperimentTracker and implement all abstract methods:

from srforge.tracking.base import ExperimentTracker

class MLflowTracker(ExperimentTracker):
    def __init__(self, experiment_name, tracking_uri=None):
        import mlflow
        mlflow.set_tracking_uri(tracking_uri)
        mlflow.set_experiment(experiment_name)
        self._run = mlflow.start_run()

    @property
    def run_id(self):
        return self._run.info.run_id

    # ... implement all other methods ...

Then use it in config:

tracker:
  _target: my_project.tracking.MLflowTracker
  params:
    experiment_name: my-experiment
    tracking_uri: http://localhost:5000

No script changes needed — the tracker is resolved from config like any other component.

Next: Hooks — Pluggable monitoring, logging, and checkpointing