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:

Calls wandb.init(**all_kwargs) on first use — the first log, or the first read of a run property such as is_resumed — not at construction.
Registers itself as the tracker the SIGINT handler in srforge.init() drains — tracker.finish(1) on Ctrl+C, so pending logs reach W&B.
Prints a Rich-styled run-info panel (Run ID / Name / Project / Path / Resumed: yes if applicable).

On Windows, W&B's helper process dies with your training

W&B uploads through a separate helper process, wandb-core. On Windows it used to outlive a crashed or force-killed training process for days, leaving the run marked "running" in W&B. W&B's own watchdog doesn't catch this on Windows, not even in its latest release (0.30.0). So just before wandb.init, the tracker ties every process started from the training process to its lifetime (a Windows job object). When training ends — however it ends — Windows stops the helper too, along with DataLoader workers and anything else the script started. If Windows refuses, a warning is logged and training continues without it. On every other system nothing changes.

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