SR-Forge

GANModel

A GANModel wraps a generator and discriminator into a single composite model for adversarial training. The generator handles the normal forward path (inference, validation), while the discriminator is orchestrated internally for training — you never call the discriminator directly.

---

Your First GANModel

from srforge.models import Model, GANModel

class Generator(Model):
    def forward(self, lrs):
        return self.net(lrs)

class Discriminator(Model):
    def forward(self, image):
        return self.disc_net(image)

gen = Generator()
gen.set_io({"inputs": {"lrs": "lrs"}, "outputs": ["sr"]})

disc = Discriminator()

gan = GANModel(generator=gen, discriminator=disc)
gan.set_io({
    "inputs": {"real_field": "hr", "fake_field": "sr"},
})

Three things to notice:

1. The generator has its own IO binding — it reads lrs from the Entry and writes sr. This binding must be set before (or during) GANModel construction.
2. The discriminator has no IO binding — GANModel orchestrates it internally.
3. GANModel's IO binding describes the discriminator's field mapping — which Entry fields are "real" (target) and "fake" (generated), and where to write the scores.

---

How It Works

GANModel has two modes:

Inference — model(entry) delegates to the generator. The discriminator is not used:

entry = Entry(lrs=input_tensor)
result = gan(entry)  # result["sr"] contains the generator output

Training — the runner calls two methods that score images through the discriminator:

entry (with hr, sr) ──> discriminator_step(entry)
                           ├── score real (hr) ──> real_score_d
                           └── score fake (sr, detached) ──> fake_score_d

                     ──> generator_step(entry)
                           ├── score real (hr, no grad) ──> real_score_g
                           └── score fake (sr, with grad) ──> fake_score_g

The D step and G step write to separate fields — all four scores coexist in the Entry without overwriting each other. This means both are available for logging and metrics.

Gradient Control

The two methods handle gradient flow automatically:

Step	Real image	Fake image
D step	Normal	.detach() — D learns, G does not
G step	torch.no_grad() — reference only	Normal — G learns through D

Why .detach() in the D step? The discriminator needs gradients for its own weights from both the real and fake scoring. .detach() on the fake tensor severs the graph at that point — D still builds a computational graph for its parameters, but gradients stop at the fake tensor and don't reach the generator.

Why torch.no_grad() in the G step? The real score is just a reference value. We don't need any gradients from it — not for D (we're updating G, not D) and not for G (the real image has no connection to the generator). no_grad skips building the graph entirely, saving memory.

Why not use no_grad in the D step too? That would prevent the discriminator from computing gradients for its own parameters — D wouldn't learn.

You don't need to manage detach or no_grad yourself — GANModel does it internally.

---

IO Binding

GANModel's set_io() configures which Entry fields the discriminator reads from and writes to.

Inputs

Two required keys:

Key	Description
real_field	Entry field containing the real (target) image
fake_field	Entry field containing the fake (generated) image

Outputs

GANModel declares four output ports via its IOSpec:

Port name	Written by	Description
real_score_d	D step	Discriminator score for real image
fake_score_d	D step	Discriminator score for (detached) fake image
real_score_g	G step	Discriminator score for real image (no grad)
fake_score_g	G step	Discriminator score for fake image (with grad)

You can inspect these programmatically:

>>> GANModel.io_spec.all_outputs
('real_score_d', 'fake_score_d', 'real_score_g', 'fake_score_g')

If you're happy with the default names, omit outputs entirely — port names are used as field names:

gan.set_io({
    "inputs": {"real_field": "hr", "fake_field": "sr"},
})

---

Adversarial Losses

SR-Forge provides two loss families. Both produce per-sample losses compatible with the standard Loss interface.

Vanilla GAN (recommended)

Standard BCE — real should score 1, fake should score 0. Each image scored independently, no coupling between real and fake:

d_criterion:
  _target: GANDiscriminatorLoss
  io: {inputs: {real_score: real_score_d, fake_score: fake_score_d}}

g_criterion:
  _target: GANGeneratorLoss
  params:
    weight: 0.005     # keep small to preserve pixel quality
  io: {inputs: {fake_score: fake_score_g}}

The G loss only needs fake_score — the generator wants D to think its fakes are real. The weight controls how much adversarial sharpening vs pixel accuracy the generator optimizes for.

Relativistic Average GAN (RaGAN)

Each score is relative to the batch average of the other. More theoretically grounded but can be less stable with spatial discriminators:

d_criterion:
  _target: RaGANDiscriminatorLoss
  io: {inputs: {real_score: real_score_d, fake_score: fake_score_d}}

g_criterion:
  _target: RaGANGeneratorLoss
  params:
    weight: 0.005
  io: {inputs: {real_score: real_score_g, fake_score: fake_score_g}}

---

Training Hooks

GANTrainingRunner uses the hook system for composable training behaviors. Instead of hardcoding R1 penalty, gradient clipping, warmup, or step scheduling in the runner, these are injected as independent hook classes via config. See Trainers & Runners → Training Hooks for the full API; the GAN-specific bits are below:

training_runner:
  _target: GANTrainingRunner
  params:
    hooks:
      - _target: StepRatio
        params: {ratio: 1.0}
      - _target: R1GradientPenalty
        params: {weight: 0.5}
      - _target: GradientClip
        params: {max_norm: 1.0, target: D}
      - _target: DWarmup
        params: {batches: 500}

Available Hooks

Hook	Stage	Description
StepRatio	pre_step	Controls G/D update frequency. ratio=2.0 trains G twice per D update.
DWarmup	pre_step	D-only training for the first N batches (across epochs).
R1GradientPenalty	post_d_forward	Penalizes large D gradients on real images. Prevents score explosion. Resolution-independent (mean reduction).
GradientClip	post_d_backward / post_g_backward	Clips gradient norms. target="D" or target="G".

Custom Hooks

Write a Hook subclass and override on_<stage> methods. They auto-bind to matching HookPoints on GANTrainingRunner at attach time; on a wrong target the Hook.bind_to call emits a DeprecationWarning.

from srforge.training.hooks import Hook
from srforge.training.runners import GANTrainingRunner

class MyHook(Hook):
    # Type-annotating the Context with the runner's nested Ctx gives
    # IDE autocomplete on the GAN-specific fields (run_d, d_extra_losses, …).
    def on_post_d_forward(self, ctx: GANTrainingRunner.Ctx):
        # Add a custom regularization term
        ctx.d_extra_losses["my_reg"] = 0.1 * some_penalty(ctx.model)

See Trainers & Runners → Training Hooks for the full HookPoint inventory, the Hook base API, and how to add new entry points in custom runners.

---

Real-World YAML Config

A complete GAN fine-tuning config based on the MagNAt_v2 super-resolution model trained on ProbaV satellite data. This is a working config — not a simplified example:

# ═══════════════════════════════════════════════════════════
# Model
# ═══════════════════════════════════════════════════════════

generator:
  _target: graph_sr.models.MISR.MagNAt.MagNAt_v2
  params:
    in_channels: 1
    scale: 3
    d: 56
    s: 16
    processing_layers: 4
    attention_heads: 1
    radius: 1.0
    edge_dropout: 0.1
    reg_feat_channels: 32
    reg_R: 4
    reg_temperature: 1.0
    align_sr_to_hr: True
  io:
    inputs: {batch: lrs}
    outputs: [sr, lr_shifts, hr_shift]

discriminator:
  _target: srforge.models.discriminator.UNetDiscriminatorSN
  params:
    num_in_ch: 2        # raw image + Laplacian edge map
    num_feat: 64
    skip_connection: True

model:
  _target: srforge.models.GANModel
  params:
    generator: ${ref:generator}
    discriminator: ${ref:discriminator}
  io:
    inputs:
      real_field: hr
      fake_field: sr

# ═══════════════════════════════════════════════════════════
# Optimizers
# ═══════════════════════════════════════════════════════════

optimizer_G:
  _target: torch.optim.Adam
  params:
    params: ${ref:generator}.trainable_params()
    lr: 1.0e-4
    betas: [0.9, 0.999]

optimizer_D:
  _target: torch.optim.Adam
  params:
    params: ${ref:discriminator}.trainable_params()
    lr: 1.0e-4
    betas: [0.0, 0.999]    # no momentum for D — standard GAN practice

# ═══════════════════════════════════════════════════════════
# Losses
# ═══════════════════════════════════════════════════════════

loss:
  _target: srforge.loss.LossCombiner
  params:
    losses:
      - _target: srforge.loss.metrics.L1
        params: {weight: 1.0}
        io: {inputs: {x: sr, y: hr}}
      - _target: graph_sr.loss.ffl.FocalFrequencyLoss
        params: {weight: 0.5, alpha: 1.0}
        io: {inputs: {x: sr, y: hr}}

d_criterion:
  _target: srforge.loss.adversarial.GANDiscriminatorLoss
  io:
    inputs: {real_score: real_score_d, fake_score: fake_score_d}

g_criterion:
  _target: srforge.loss.adversarial.GANGeneratorLoss
  params:
    weight: 0.005               # keep small — pixel quality first
  io:
    inputs: {fake_score: fake_score_g}

# ═══════════════════════════════════════════════════════════
# Runner
# ═══════════════════════════════════════════════════════════

training_runner:
  _target: srforge.training.runners.GANTrainingRunner
  params:
    optimizer_G: ${ref:optimizer_G}
    optimizer_D: ${ref:optimizer_D}
    d_criterion: ${ref:d_criterion}
    g_criterion: ${ref:g_criterion}
    hooks:
      - _target: srforge.training.hooks.StepRatio
        params: {ratio: 1.0}
      - _target: srforge.training.hooks.EdgeEnhancedInput
    device: ${system.device}
    empty_cache_every: 4

trainer:
  _target: srforge.training.trainers.PyTorchTrainer
  params:
    model: ${ref:model}
    training_epoch_runner: ${ref:training_runner}
    validation_epoch_runner: ${ref:validation_runner}
    training_criterion: ${ref:loss}
    validation_criterion: ${ref:validation_metrics}
    lr_scheduler: ${ref:lr_scheduler}

Key Design Decisions

Generator weight (g_criterion.weight: 0.005) — The adversarial contribution should be roughly equal to or smaller than the pixel losses. Too high (0.1+) causes PSNR to collapse as G prioritizes fooling D over pixel accuracy. Too low (0.0001) makes the adversarial signal negligible. Start at 0.005 and tune based on validation PSNR.

D learning rate and momentum (betas: [0.0, 0.999]) — Zero momentum (beta1=0) for the discriminator is standard GAN practice. It prevents D from building up aggressive momentum that causes score explosion. D learning rate should be comparable to G (1e-4) — if too high, D dominates; if too low, D can't learn.

Edge-enhanced discriminator — For single-channel or low-contrast images, the discriminator may struggle to detect subtle quality differences. Concatenating a Laplacian edge map amplifies high-frequency differences (blur vs sharp), enabling D to learn on data where raw pixel differences are tiny.

No R1 penalty by default — R1 gradient penalty (penalizes large D gradients on real images) prevents score explosion but can suppress D learning if too strong. With vanilla GAN loss (which is inherently more stable than RaGAN), R1 may not be needed. Add it as a hook if D loss explodes.

Training Dynamics to Watch

At equilibrium, both D loss and G loss should hover around 0.4-0.7 (for vanilla GAN with per-sample BCE). Watch for:

Pattern	Meaning	Action
Both ~0.5-0.7, oscillating	Healthy dynamic equilibrium	Let it train
D dropping, G rising	D winning	Lower D LR or increase G weight
D rising, G dropping	G winning	Rare with pre-trained G — increase D LR
Both at exactly 0.693	Dead equilibrium — D outputs zeros	Remove R1 or increase D LR
D > 3 and climbing	Score explosion	Add R1 hook or lower D LR
PSNR dropping steadily	Adversarial weight too high	Lower G weight

---

Programmatic Use

For training loops outside the config system:

import torch
from srforge.models import Model, GANModel
from srforge.data import Entry

# 1. Build models
gen = MyGenerator()
gen.set_io({"inputs": {"image": "lrs"}, "outputs": ["sr"]})

disc = MyDiscriminator()

gan = GANModel(generator=gen, discriminator=disc)
gan.set_io({"inputs": {"real_field": "hr", "fake_field": "sr"}})

# 2. Forward pass
entry = Entry(lrs=lr_tensor, hr=hr_tensor)
entry = gan(entry)             # generator forward: produces sr

# 3. D step
entry = gan.discriminator_step(entry)
# entry now has: real_score_d, fake_score_d

# 4. G step
entry = gan.generator_step(entry)
# entry now has: real_score_g, fake_score_g (plus D-step scores)

trainable_params()

Returns only the generator's trainable parameters. The discriminator gets its own optimizer via discriminator.trainable_params():

optimizer_G = torch.optim.Adam(gan.trainable_params(), lr=1e-4)
optimizer_D = torch.optim.Adam(gan.discriminator.trainable_params(), lr=1e-4)

---

IOSpec and Discoverability

GANModel declares a class-level IOSpec, making its interface discoverable:

>>> from srforge.models import GANModel
>>> GANModel.io_spec.all_inputs
('real_field', 'fake_field')
>>> GANModel.io_spec.all_outputs
('real_score_d', 'fake_score_d', 'real_score_g', 'fake_score_g')

The output dict keys in set_io() are validated against this spec. Unknown keys raise ValueError:

gan.set_io({
    "inputs": {"real_field": "hr", "fake_field": "sr"},
    "outputs": {"typo_field": "x"},
})
# ValueError: GANModel outputs has unknown port names: ['typo_field'].
#   Valid: ['fake_score_d', 'fake_score_g', 'real_score_d', 'real_score_g'].

---

Next: SequentialModel — Compose models and transforms into multi-stage pipelines