KLGaussian

class colibri.regularizers.KLGaussian(mean=0.01, stddev=2.0)

Bases: Module

KL Divergence Regularization for Gaussian Distributions.

Code adapted from [2] Jacome, Roman, Pablo Gomez, and Henry Arguello. “Middle output regularized end-to-end optimization for computational imaging.” Optica 10.11 (2023): 1421-1431.

$$R(\mathbf{y})=\text{KL}(p_{\mathbf{y}},p_{\mathbf{z}})=-\frac{1}{2}\sum _{i=1}^n(1+\log ({\sigma }_{{\mathbf{y}}_i}^2)-\log ({\sigma }_{{\mathbf{z}}_i}^2)-\frac{{\sigma }_{{\mathbf{y}}_i}^2+({\mu }_{{\mathbf{y}}_i}-{\mu }_{{\mathbf{z}}_i}{)}^2}{{\sigma }_{{\mathbf{z}}_i}^2})$$

where ${\mu }_{{\mathbf{y}}_i}$ and ${\sigma }_{{\mathbf{y}}_i}$ are the mean and standard deviation of the input tensor $\mathbf{y}\in \mathcal{Y}$, respectively, and ${\mu }_{{\mathbf{z}}_i}$ and ${\sigma }_{{\mathbf{z}}_i}$ are the target mean and standard deviation, respectively.

forward(y)

Compute KL divergence regularization term.

Parameters:

y (torch.Tensor) – Input tensor representing a Gaussian distribution.

Returns:

KL divergence regularization term.

Return type:

torch.Tensor

Examples using KLGaussian:

Demo Colibri.

Demo Colibri.