EquivariantDenoiser

class deepinv.models.EquivariantDenoiser(denoiser, transform: Transform | None = None, random=True)[source]

Bases: Module

Turns the input denoiser into an equivariant denoiser with respect to geometric transforms.

Recall that a denoiser is equivariant with respect to a group of transformations if it commutes with the action of the group. More precisely, let \(\mathcal{G}\) be a group of transformations \(\{T_g\}_{g\in \mathcal{G}}\) and \(\denoisername\) a denoiser. Then, \(\denoisername\) is equivariant with respect to \(\mathcal{G}\) if \(\denoisername(T_g(x)) = T_g(\denoisername(x))\) for any image \(x\) and any \(g\in \mathcal{G}\).

The denoiser can be turned into an equivariant denoiser by averaging over the group of transforms, i.e.

\[\operatorname{D}^{\text{eq}}_{\sigma}(x) = \frac{1}{|\mathcal{G}|}\sum_{g\in \mathcal{G}} T_g^{-1}(\operatorname{D}_{\sigma}(T_g(x))).\]

Otherwise, as proposed in https://arxiv.org/abs/2312.01831, a Monte Carlo approximation can be obtained by sampling \(g \sim \mathcal{G}\) at random and applying

\[\operatorname{D}^{\text{MC}}_{\sigma}(x) = T_g^{-1}(\operatorname{D}_{\sigma}(T_g(x))).\]

Note

We have implemented many popular geometric transforms, see docs. You can set the number of Monte Carlo samples by passing n_trans into the transforms, for example Rotate(n_trans=2) will average over 2 samples per call. For rotate and reflect, by setting n_trans to the maximum (e.g. 4 for 90 degree rotations, 2 for 1D reflections), it will average over the whole group, for example:

Rotate(n_trans=4, multiples=90, positive=True) * Reflect(n_trans=2, dims=[-1])

See Image transforms for equivariance & augmentations for an example.

Parameters:

  • denoiser (callable) – Denoiser \(\operatorname{D}_{\sigma}\).

  • transform (Transform) – geometric transformation. If None, defaults to rotations of multiples of 90 with horizontal flips (see note above). See docs for list of available transforms.

  • random (bool) – if True, the denoiser is applied to a randomly transformed version of the input image each time i.e. a Monte-Carlo approximation of an equivariant denoiser. If False, the denoiser is applied to the average of all the transformed images, turning the denoiser into an equivariant denoiser with respect to the chosen group of transformations. Ignored if transform is provided.

forward(x, *denoiser_args, **denoiser_kwargs)[source]

Symmetrize the denoiser by the transformation to create an equivariant denoiser and apply to input.

The symmetrization collects the average if multiple samples are used (controlled with n_trans in the transform).

Parameters:

  • x (torch.Tensor) – input image.

  • *denoiser_args – args for denoiser function e.g. sigma noise level.

  • **denoiser_kwargs – kwargs for denoiser function e.g. sigma noise level.

Returns:

denoised image.

Examples using EquivariantDenoiser:

Image transforms for equivariance & augmentations

Image transforms for equivariance & augmentations