LION.models.PnP.gradient_step_denoiser

Classes

GSD([model_parameters, denoiser])

The gradient step denoiser: x - grad g(x), with g(x) = 0.5 ||x - N(x)||^2 for a denoiser N (Hurault et al., 2022).

---

class LION.models.PnP.gradient_step_denoiser.GSD(model_parameters=None, denoiser=None)

Bases: LIONmodel

The gradient step denoiser: x - grad g(x), with g(x) = 0.5 ||x - N(x)||^2 for a denoiser N (Hurault et al., 2022).

denoiser: N, any LIONmodel (to use a plain torch module, wrap it in a LIONmodel, e.g. TorchModel). Without it, N is the CNN model_parameters.backbone_model (default DRUNet) made with model_parameters.backbone_params.

The noise level: GSD takes one if N does (N.takes_noise_level()), and passes it on; how N uses it (e.g. DRUNet’s extra input channel) is N’s.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

spec_kwargs()

GSD’s constructor arguments: its parameters, and the denoiser when one was given (it is saved inside GSD’s spec).

static default_parameters()

backbone_model (the denoiser N’s class, DRUNet) and backbone_params (its parameters: a small DRUNet of 32 channels and 2 blocks per scale, without the noise level). Ignored when GSD is given a denoiser.

static cite(cite_format='bib')

Where GSD comes from: its citation, BibTeX (“bib”) or “MLA”.

grad(x, noise_level=None)

g(x) = 0.5 ||x - N(x)||^2 and its gradient with respect to x.

Differentiable with respect to the weights in training, and with respect to x when x requires a gradient (e.g. GSD inside a network trained end to end).

takes_noise_level()

Whether GSD takes a noise level: if its denoiser does.

forward(x, noise_level=None)

The denoised images: x - grad g(x) (grad), given noise_level if the denoiser takes one.