deepinv 0.1.0.dev0__py3-none-any.whl

deepinv/loss/sure.py

@@ -0,0 +1,338 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def hutch_div(y, physics, f, mc_iter=1):
+    r"""
+    Hutch divergence for A(f(x)).
+
+    :param torch.Tensor y: Measurements.
+    :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+    :param torch.nn.Module f: Reconstruction network.
+    :param int mc_iter: number of iterations. Default=1.
+    :return: (float) hutch divergence.
+    """
+    input = y.requires_grad_(True)
+    output = physics.A(f(input, physics))
+    out = 0
+    for i in range(mc_iter):
+        b = torch.randn_like(input)
+        x = torch.autograd.grad(output, input, b, retain_graph=True, create_graph=True)[
+            0
+        ]
+        out += (b * x).mean()
+
+    return out / mc_iter
+
+
+def exact_div(y, physics, model):
+    r"""
+    Exact divergence for A(f(x)).
+
+    :param torch.Tensor y: Measurements.
+    :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+    :param torch.nn.Module model: Reconstruction network.
+    :param int mc_iter: number of iterations. Default=1.
+    :return: (float) exact divergence.
+    """
+    input = y.requires_grad_(True)
+    output = physics.A(model(input, physics))
+    out = 0
+    _, c, h, w = input.shape
+    for i in range(c):
+        for j in range(h):
+            for k in range(w):
+                b = torch.zeros_like(input)
+                b[:, i, j, k] = 1
+                x = torch.autograd.grad(
+                    output, input, b, retain_graph=True, create_graph=True
+                )[0]
+                out += (b * x).sum()
+
+    return out / (c * h * w)
+
+
+def mc_div(y1, y, f, physics, tau):
+    r"""
+    Monte-Carlo estimation for the divergence of A(f(x)).
+
+    :param torch.Tensor y: Measurements.
+    :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+    :param torch.nn.Module f: Reconstruction network.
+    :param int mc_iter: number of iterations. Default=1.
+    :return: (float) hutch divergence.
+    """
+    b = torch.randn_like(y)
+    y2 = physics.A(f(y + b * tau, physics))
+    out = (b * (y2 - y1) / tau).mean()
+    return out
+
+
+class SureGaussianLoss(nn.Module):
+    r"""
+    SURE loss for Gaussian noise
+
+    The loss is designed for the following noise model:
+
+    .. math::
+
+        y \sim\mathcal{N}(u,\sigma^2 I) \quad \text{with}\quad u= A(x).
+
+    The loss is computed as
+
+    .. math::
+
+        \frac{1}{m}\|y - A\inverse{y}\|_2^2 -\sigma^2 +\frac{2\sigma^2}{m\tau}b^{\top} \left(A\inverse{y+\tau b_i} -
+        A\inverse{y}\right)
+
+    where :math:`R` is the trainable network, :math:`A` is the forward operator,
+    :math:`y` is the noisy measurement vector of size :math:`m`, :math:`A` is the forward operator,
+    :math:`b\sim\mathcal{N}(0,I)` and :math:`\tau\geq 0` is a hyperparameter controlling the
+    Monte Carlo approximation of the divergence.
+
+    This loss approximates the divergence of :math:`A\inverse{y}` (in the original SURE loss)
+    using the Monte Carlo approximation in
+    https://ieeexplore.ieee.org/abstract/document/4099398/
+
+    If the measurement data is truly Gaussian with standard deviation :math:`\sigma`,
+    this loss is an unbiased estimator of the mean squared loss :math:`\frac{1}{m}\|u-A\inverse{y}\|_2^2`
+    where :math:`z` is the noiseless measurement.
+
+    .. warning::
+
+        The loss can be sensitive to the choice of :math:`\tau`, which should be proportional to the size of :math:`y`.
+        The default value of 0.01 is adapted to :math:`y` vectors with entries in :math:`[0,1]`.
+
+    :param float sigma: Standard deviation of the Gaussian noise.
+    :param float tau: Approximation constant for the Monte Carlo approximation of the divergence.
+    """
+
+    def __init__(self, sigma, tau=1e-2):
+        super(SureGaussianLoss, self).__init__()
+        self.name = "SureGaussian"
+        self.sigma2 = sigma**2
+        self.tau = tau
+
+    def forward(self, y, x_net, physics, model, **kwargs):
+        r"""
+        Computes the SURE Loss.
+
+        :param torch.Tensor y: Measurements.
+        :param torch.Tensor x_net: reconstructed image :math:`\inverse{y}`.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements.
+        :param torch.nn.Module model: Reconstruction network.
+        :return: (float) SURE loss.
+        """
+
+        y1 = physics.A(x_net)
+        div = 2 * self.sigma2 * mc_div(y1, y, model, physics, self.tau)
+        mse = (y1 - y).pow(2).mean()
+        loss_sure = mse + div - self.sigma2
+        return loss_sure
+
+
+class SurePoissonLoss(nn.Module):
+    r"""
+    SURE loss for Poisson noise
+
+    The loss is designed for the following noise model:
+
+    .. math::
+
+      y = \gamma z \quad \text{with}\quad z\sim \mathcal{P}(\frac{u}{\gamma}), \quad u=A(x).
+
+    The loss is computed as
+
+    .. math::
+
+        \frac{1}{m}\|y-A\inverse{y}\|_2^2-\frac{\gamma}{m} 1^{\top}y
+        +\frac{2\gamma}{m\tau}(b\odot y)^{\top} \left(A\inverse{y+\tau b}-A\inverse{y}\right)
+
+    where :math:`R` is the trainable network, :math:`y` is the noisy measurement vector of size :math:`m`,
+    :math:`b` is a Bernoulli random variable taking values of -1 and 1 each with a probability of 0.5,
+    :math:`\tau` is a small positive number, and :math:`\odot` is an elementwise multiplication.
+
+    See https://ieeexplore.ieee.org/abstract/document/6714502/ for details.
+    If the measurement data is truly Poisson
+    this loss is an unbiased estimator of the mean squared loss :math:`\frac{1}{m}\|u-A\inverse{y}\|_2^2`
+    where :math:`z` is the noiseless measurement.
+
+    .. warning::
+
+        The loss can be sensitive to the choice of :math:`\tau`, which should be proportional to the size of :math:`y`.
+        The default value of 0.01 is adapted to :math:`y` vectors with entries in :math:`[0,1]`.
+
+    :param float gain: Gain of the Poisson Noise.
+    :param float tau: Approximation constant for the Monte Carlo approximation of the divergence.
+    """
+
+    def __init__(self, gain, tau=1e-3):
+        super(SurePoissonLoss, self).__init__()
+        self.name = "SurePoisson"
+        self.gain = gain
+        self.tau = tau
+
+    def forward(self, y, x_net, physics, model, **kwargs):
+        r"""
+        Computes the SURE loss.
+
+        :param torch.Tensor y: measurements.
+        :param torch.Tensor x_net: reconstructed image :math:`\inverse{y}`.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements
+        :param torch.nn.Module model: Reconstruction network
+        :return: (float) SURE loss.
+        """
+
+        # generate a random vector b
+        b = torch.rand_like(y) > 0.5
+        b = (2 * b - 1) * 1.0  # binary [-1, 1]
+
+        y1 = physics.A(x_net)
+        y2 = physics.A(model(y + self.tau * b, physics))
+
+        # compute m (size of y)
+        # m = y.numel() #(torch.abs(y) > 1e-5).flatten().sum()
+
+        loss_sure = (
+            (y1 - y).pow(2).mean()
+            - self.gain * y.mean()
+            + 2.0 / self.tau * (b * y * self.gain * (y2 - y1)).mean()
+        )
+
+        return loss_sure
+
+
+class SurePGLoss(nn.Module):
+    r"""
+    SURE loss for Poisson-Gaussian noise
+
+    The loss is designed for the following noise model:
+
+    .. math::
+
+        y = \gamma z + \epsilon
+
+    where :math:`u = A(x)`, :math:`z \sim \mathcal{P}\left(\frac{u}{\gamma}\right)`,
+    and :math:`\epsilon \sim \mathcal{N}(0, \sigma^2 I)`.
+
+    The loss is computed as
+
+    .. math::
+
+        & \frac{1}{m}\|y-A\inverse{y}\|_2^2-\frac{\gamma}{m} 1^{\top}y-\sigma^2
+        +\frac{2}{m\tau_1}(b\odot (\gamma y + \sigma^2 I))^{\top} \left(A\inverse{y+\tau b}-A\inverse{y} \right) \\\\
+        & +\frac{2\gamma \sigma^2}{m\tau_2^2}c^{\top} \left( A\inverse{y+\tau c} + A\inverse{y-\tau c} - 2A\inverse{y} \right)
+
+    where :math:`R` is the trainable network, :math:`y` is the noisy measurement vector,
+    :math:`b` is a Bernoulli random variable taking values of -1 and 1 each with a probability of 0.5,
+    :math:`\tau` is a small positive number, and :math:`\odot` is an elementwise multiplication.
+
+    If the measurement data is truly Poisson-Gaussian
+    this loss is an unbiased estimator of the mean squared loss :math:`\frac{1}{m}\|u-A\inverse{y}\|_2^2`
+    where :math:`z` is the noiseless measurement.
+
+    See https://ieeexplore.ieee.org/abstract/document/6714502/ for details.
+
+    .. warning::
+
+        The loss can be sensitive to the choice of :math:`\tau`, which should be proportional to the size of :math:`y`.
+        The default value of 0.01 is adapted to :math:`y` vectors with entries in :math:`[0,1]`.
+
+    :param float sigma: Standard deviation of the Gaussian noise.
+    :param float gamma: Gain of the Poisson Noise.
+    :param float tau: Approximation constant for the Monte Carlo approximation of the divergence.
+    """
+
+    def __init__(self, sigma, gain, tau1=1e-3, tau2=1e-2):
+        super(SurePGLoss, self).__init__()
+        self.name = "SurePG"
+        # self.sure_loss_weight = sure_loss_weight
+        self.sigma2 = sigma**2
+        self.gain = gain
+        self.tau1 = tau1
+        self.tau2 = tau2
+
+    def forward(self, y, x_net, physics, model, **kwargs):
+        r"""
+        Computes the SURE loss.
+
+        :param torch.Tensor y: measurements.
+        :param torch.Tensor x_net: reconstructed image :math:`\inverse{y}`.
+        :param deepinv.physics.Physics physics: Forward operator associated with the measurements
+        :param torch.nn.Module f: Reconstruction network
+        :return: (float) SURE loss.
+        """
+
+        b1 = torch.rand_like(y) > 0.5
+        b1 = (2 * b1 - 1) * 1.0  # binary [-1, 1]
+
+        p = 0.7236  # .5 + .5*np.sqrt(1/5.)
+
+        b2 = torch.ones_like(b1) * np.sqrt(p / (1 - p))
+        b2[torch.rand_like(b2) < p] = -np.sqrt((1 - p) / p)
+
+        meas1 = physics.A(x_net)
+        meas2 = physics.A(model(y + self.tau1 * b1, physics))
+        meas2p = physics.A(model(y + self.tau2 * b2, physics))
+        meas2n = physics.A(model(y - self.tau2 * b2, physics))
+
+        # compute m (size of y)
+        # m = (torch.abs(y) > 1e-5).flatten().sum()
+
+        loss_mc = (meas1 - y).pow(2).mean()
+
+        loss_div1 = (
+            2
+            / self.tau1
+            * ((b1 * (self.gain * y + self.sigma2)) * (meas2 - meas1)).mean()
+        )
+
+        offset = -self.gain * y.mean() - self.sigma2
+
+        loss_div2 = (
+            -2
+            * self.sigma2
+            * self.gain
+            / (self.tau2**2)
+            * (b2 * (meas2p + meas2n - 2 * meas1)).mean()
+        )
+
+        loss_sure = loss_mc + loss_div1 + loss_div2 + offset
+        return loss_sure
+
+
+# if __name__ == "__main__":
+#     from deepinv.models import Denoiser
+#     import deepinv as dinv
+#
+#     model_spec = {
+#         "name": "waveletprior",
+#         "args": {"wv": "db8", "level": 3, "device": dinv.device},
+#     }
+#     f = dinv.models.ArtifactRemoval(Denoiser(model_spec))
+#     # test divergence
+#
+#     x = torch.ones((1, 3, 16, 16), device=dinv.device) * 0.5
+#     physics = dinv.physics.Denoising(dinv.physics.GaussianNoise(0.1))
+#     y = physics(x)
+#
+#     y1 = f(y, physics)
+#     tau = 1e-4
+#
+#     exact = exact_div(y, physics, f)
+#
+#     error_h = 0
+#     error_mc = 0
+#     for i in range(100):
+#         h = hutch_div(y, physics, f)
+#         mc = mc_div(y1, y, f, physics, tau)
+#
+#         error_h += torch.abs(h - exact)
+#         error_mc += torch.abs(mc - exact)
+#
+#     error_mc /= 100
+#     error_h /= 100
+#
+#     print(f"error_h: {error_h}")
+#     print(f"error_mc: {error_mc}")

deepinv/loss/tv.py

@@ -0,0 +1,39 @@
+import torch
+import torch.nn as nn
+
+
+class TVLoss(nn.Module):
+    r"""
+    Total variation loss (:math:`\ell_2` norm).
+
+    It computes the loss :math:`\|D\hat{x}\|_2^2`,
+    where :math:`D` is a normalized linear operator that computes the vertical and horizontal first order differences
+    of the reconstructed image :math:`\hat{x}`.
+
+    :param float weight: scalar weight for the TV loss.
+    """
+
+    def __init__(self, weight=1.0):
+        super(TVLoss, self).__init__()
+        self.tv_loss_weight = weight
+        self.name = "tv"
+
+    def forward(self, x_net, **kwargs):
+        r"""
+        Computes the TV loss.
+
+        :param torch.Tensor x_net: reconstructed image.
+        :return: (torch.Tensor) loss.
+        """
+        batch_size = x_net.size()[0]
+        h_x = x_net.size()[2]
+        w_x = x_net.size()[3]
+        count_h = self.tensor_size(x_net[:, :, 1:, :])
+        count_w = self.tensor_size(x_net[:, :, :, 1:])
+        h_tv = torch.pow((x_net[:, :, 1:, :] - x_net[:, :, : h_x - 1, :]), 2).sum()
+        w_tv = torch.pow((x_net[:, :, :, 1:] - x_net[:, :, :, : w_x - 1]), 2).sum()
+        return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size
+
+    @staticmethod
+    def tensor_size(t):
+        return t.size()[1] * t.size()[2] * t.size()[3]

deepinv/models/GSPnP.py

@@ -0,0 +1,129 @@
+import torch
+import torch.nn as nn
+from .utils import get_weights_url
+
+
+class StudentGrad(nn.Module):
+    def __init__(self, denoiser):
+        super().__init__()
+        self.model = denoiser
+
+    def forward(self, x, sigma):
+        return self.model(x, sigma)
+
+
+class GSPnP(nn.Module):
+    r"""
+    Gradient Step module to use a denoiser architecture as a Gradient Step Denoiser.
+    See https://arxiv.org/pdf/2110.03220.pdf.
+    Code from https://github.com/samuro95/GSPnP.
+
+    :param nn.Module denoiser: Denoiser model.
+    :param float alpha: Relaxation parameter
+    """
+
+    def __init__(self, denoiser, alpha=1.0, train=False):
+        super().__init__()
+        self.student_grad = StudentGrad(denoiser)
+        self.alpha = alpha
+        self.train = train
+
+    def potential(self, x, sigma):
+        N = self.student_grad(x, sigma)
+        return (
+            0.5
+            * self.alpha
+            * torch.norm((x - N).view(x.shape[0], -1), p=2, dim=-1) ** 2
+        )
+
+    def potential_grad(self, x, sigma):
+        r"""
+        Calculate :math:`\nabla g` the gradient of the regularizer :math:`g` at input :math:`x`.
+
+        :param torch.tensor x: Input image
+        :param float sigma: Denoiser level :math:`\sigma` (std)
+        """
+        torch.set_grad_enabled(True)
+        x = x.float()
+        x = x.requires_grad_()
+        N = self.student_grad(x, sigma)
+        JN = torch.autograd.grad(
+            N, x, grad_outputs=x - N, create_graph=True, only_inputs=True
+        )[0]
+        if not self.train:
+            torch.set_grad_enabled(False)
+        Dg = x - N - JN
+        return self.alpha * Dg
+
+    def forward(self, x, sigma):
+        r"""
+        Denoising with Gradient Step Denoiser
+
+        :param torch.tensor x: Input image
+        :param float sigma: Denoiser level (std)
+        """
+        Dg = self.potential_grad(x, sigma)
+        x_hat = x - Dg
+        return x_hat
+
+
+def GSDRUNet(
+    alpha=1.0,
+    in_channels=3,
+    out_channels=3,
+    nb=2,
+    nc=[64, 128, 256, 512],
+    act_mode="E",
+    pretrained=None,
+    train=False,
+    device=torch.device("cpu"),
+):
+    """
+    Gradient Step Denoiser with DRUNet architecture
+
+    :param float alpha: Relaxation parameter
+    :param int in_channels: Number of input channels
+    :param int out_channels: Number of output channels
+    :param int nb: Number of blocks in the DRUNet
+    :param list nc: Number of channels in the DRUNet
+    :param str act_mode: activation mode, "R" for ReLU, "L" for LeakyReLU "E" for ELU and "S" for Softplus.
+    :param str downsample_mode: Downsampling mode, "avgpool" for average pooling, "maxpool" for max pooling, and
+        "strideconv" for convolution with stride 2.
+    :param str upsample_mode: Upsampling mode, "convtranspose" for convolution transpose, "pixelsuffle" for pixel
+        shuffling, and "upconv" for nearest neighbour upsampling with additional convolution.
+    :param bool download: use a pretrained network. If ``pretrained=None``, the weights will be initialized at random
+        using Pytorch's default initialization. If ``pretrained='download'``, the weights will be downloaded from an
+        online repository (only available for the default architecture).
+        Finally, ``pretrained`` can also be set as a path to the user's own pretrained weights.
+        See :ref:`pretrained-weights <pretrained-weights>` for more details.
+    :param bool train: training or testing mode.
+    :param str device: gpu or cpu.
+
+    """
+    from deepinv.models.drunet import DRUNet
+
+    denoiser = DRUNet(
+        in_channels=in_channels,
+        out_channels=out_channels,
+        nb=nb,
+        nc=nc,
+        act_mode=act_mode,
+        pretrained=None,
+        train=train,
+        device=device,
+    )
+    GSmodel = GSPnP(denoiser, alpha=alpha, train=train)
+    if pretrained:
+        if pretrained == "download":
+            url = get_weights_url(model_name="gradientstep", file_name="GSDRUNet.ckpt")
+            ckpt = torch.hub.load_state_dict_from_url(
+                url,
+                map_location=lambda storage, loc: storage,
+                file_name="GSDRUNet.ckpt",
+            )["state_dict"]
+        else:
+            ckpt = torch.load(pretrained, map_location=lambda storage, loc: storage)[
+                "state_dict"
+            ]
+        GSmodel.load_state_dict(ckpt, strict=False)
+    return GSmodel

deepinv/models/PDNet.py

@@ -0,0 +1,109 @@
+# This is an implementation of https://arxiv.org/abs/1707.06474
+import torch
+import torch.nn as nn
+
+
+def init_weights(m):
+    if isinstance(m, nn.Linear):
+        torch.nn.init.xavier_uniform(m.weight)
+        m.bias.data.fill_(0.0)
+
+
+class PDNet_PrimalBlock(nn.Module):
+    def __init__(self, in_channels=6, out_channels=5, depth=3, bias=True, nf=32):
+        r"""
+        Primal block for the Primal-Dual unfolding model (PDNet) from https://arxiv.org/abs/1707.06474.
+
+        Primal variables are images of shape (batch_size, in_channels, height, width). The input of each
+        primal block is the concatenation of the current primal variable and the backprojected dual variable along
+        the channel dimension. The output of each primal block is the current primal variable.
+
+        :param int in_channels: number of input channels. Default: 6.
+        :param int out_channels: number of output channels. Default: 5.
+        :param int depth: number of convolutional layers in the block. Default: 3.
+        :param bool bias: whether to use bias in convolutional layers. Default: True.
+        :param int nf: number of features in the convolutional layers. Default: 32.
+        """
+        super(PDNet_PrimalBlock, self).__init__()
+
+        self.depth = depth
+
+        self.in_conv = nn.Conv2d(
+            in_channels, nf, kernel_size=3, stride=1, padding=1, bias=bias
+        )
+        self.in_conv.apply(init_weights)
+        self.conv_list = nn.ModuleList(
+            [
+                nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1, bias=bias)
+                for _ in range(self.depth - 2)
+            ]
+        )
+        self.conv_list.apply(init_weights)
+        self.out_conv = nn.Conv2d(
+            nf, out_channels, kernel_size=3, stride=1, padding=1, bias=bias
+        )
+        self.out_conv.apply(init_weights)
+
+        self.nl_list = nn.ModuleList([nn.PReLU() for _ in range(self.depth - 1)])
+
+    def forward(self, x, Atu):
+        x_in = torch.cat((x, Atu), dim=1)
+
+        x_ = self.in_conv(x_in)
+        x_ = self.nl_list[0](x_)
+
+        for i in range(self.depth - 2):
+            x_l = self.conv_list[i](x_)
+            x_ = self.nl_list[i + 1](x_l)
+
+        return self.out_conv(x_) + x
+
+
+class PDNet_DualBlock(nn.Module):
+    def __init__(self, in_channels=7, out_channels=5, depth=3, bias=True, nf=32):
+        r"""
+        Dual block for the Primal-Dual unfolding model (PDNet) from https://arxiv.org/abs/1707.06474.
+
+        Dual variables are images of shape (batch_size, in_channels, height, width). The input of each
+        primal block is the concatenation of the current dual variable with the projected primal variable and
+        the measurements. The output of each dual block is the current primal variable.
+
+        :param int in_channels: number of input channels. Default: 7.
+        :param int out_channels: number of output channels. Default: 5.
+        :param int depth: number of convolutional layers in the block. Default: 3.
+        :param bool bias: whether to use bias in convolutional layers. Default: True.
+        :param int nf: number of features in the convolutional layers. Default: 32.
+        """
+        super(PDNet_DualBlock, self).__init__()
+
+        self.depth = depth
+
+        self.in_conv = nn.Conv2d(
+            in_channels, nf, kernel_size=3, stride=1, padding=1, bias=bias
+        )
+        self.in_conv.apply(init_weights)
+        self.conv_list = nn.ModuleList(
+            [
+                nn.Conv2d(nf, nf, kernel_size=3, stride=1, padding=1, bias=bias)
+                for _ in range(self.depth - 2)
+            ]
+        )
+        self.conv_list.apply(init_weights)
+        self.out_conv = nn.Conv2d(
+            nf, out_channels, kernel_size=3, stride=1, padding=1, bias=bias
+        )
+        self.out_conv.apply(init_weights)
+
+        self.nl_list = nn.ModuleList([nn.PReLU() for _ in range(self.depth - 1)])
+
+    def forward(self, u, Ax_cur, y):
+        x_in = torch.cat((u, Ax_cur, y), dim=1)
+
+        x_ = self.in_conv(x_in)
+        x_ = self.nl_list[0](x_)
+
+        for i in range(self.depth - 2):
+            x_l = self.conv_list[i](x_)
+            x_ = self.nl_list[i + 1](x_l)
+
+        return self.out_conv(x_) + u

deepinv/models/__init__.py

@@ -0,0 +1,17 @@
+from .drunet import DRUNet
+from .scunet import SCUNet
+from .ae import AutoEncoder
+from .unet import UNet
+from .dncnn import DnCNN
+from .artifactremoval import ArtifactRemoval
+from .tgv import TGV as TGV
+from .tv import TV as TV
+from .wavdict import WaveletPrior, WaveletDict
+from .GSPnP import GSDRUNet
+from .median import MedianFilter
+from .dip import DeepImagePrior, ConvDecoder
+from .diffunet import DiffUNet
+from .swinir import SwinIR
+from .PDNet import PDNet_PrimalBlock, PDNet_DualBlock
+from .bm3d import BM3D
+from .equivariant import EquivariantDenoiser

deepinv/models/ae.py

@@ -0,0 +1,43 @@
+import torch
+
+
+class AutoEncoder(torch.nn.Module):
+    r"""
+    Simple fully connected autoencoder network.
+
+    Simple architecture that can be used for debugging or fast prototyping.
+
+    :param int dim_input: total number of elements (pixels) of the input.
+    :param int dim_hid: number of features in intermediate layer.
+    :param int dim_hid: latent space dimension.
+    :param int residual: use a residual connection between input and output.
+
+    """
+
+    def __init__(self, dim_input, dim_mid=1000, dim_hid=32, residual=True):
+        super().__init__()
+        self.residual = residual
+
+        self.encoder = torch.nn.Sequential(
+            torch.nn.Linear(dim_input, dim_mid),
+            torch.nn.ReLU(),
+            torch.nn.Linear(dim_mid, dim_hid),
+        )
+        self.decoder = torch.nn.Sequential(
+            torch.nn.Linear(dim_hid, dim_mid),
+            torch.nn.ReLU(),
+            torch.nn.Linear(dim_mid, dim_input),
+        )
+
+    def forward(self, x, sigma=None):
+        N, C, H, W = x.shape
+        x = x.view(N, -1)
+
+        encoded = self.encoder(x)
+        decoded = self.decoder(encoded)
+
+        if self.residual:
+            decoded = decoded + x
+
+        decoded = decoded.view(N, C, H, W)
+        return decoded