UNet#
- class deepinv.models.UNet(in_channels=1, out_channels=1, residual=True, circular_padding=False, cat=True, bias=True, batch_norm=True, scales=None, channels_per_scale=None, device='cpu', dim=2)[source]#
Bases:
DenoiserU-Net convolutional denoiser.
This network is a fully convolutional denoiser based on the U-Net architecture. The architecture follows the design described in Jin et al.[1], which is adapted for inverse problems by using zero padding and a global residual connection. This differs from the original U-Net [2] commonly used for medical segmentation.
The number of stages in the network is controlled by
scales. The width of each stage is controlled bychannels_per_scale, which gives the number of feature maps at each stage, from shallow to deeper stages.The number of trainable parameters increases with both
scalesand the values inchannels_per_scale. Ifscalesis not given, it is inferred fromchannels_per_scale. If both are omitted, defaults tochannels_per_scale=[64, 128, 256, 512]. If onlyscalesis specified,channels_per_scale=[64 * (2**k) for k in range(scales)]. When both are specified,scalesmust match the length ofchannels_per_scale.Warning
When using the bias-free batch norm via
batch_norm="biasfree", NaNs may be encountered during training, causing the whole training procedure to fail.- Parameters:
in_channels (int) – input image channels
out_channels (int) – output image channels
residual (bool) – use a skip-connection between input and output.
circular_padding (bool) – circular padding for the convolutional layers.
cat (bool) – use skip-connections between intermediate levels.
bias (bool) – use learnable biases in conv and norm layers.
batch_norm (bool, str) – if False, disable normalization entirely, if
True, use batch normalization, ifbatch_norm="biasfree", use the bias-free batch norm from Mohan et al.[3].scales (int) – Number of stages.
channels_per_scale (Sequence[int]) – Number of feature maps at each stage (from shallow to deep).
device (torch.device, str) – Device to put the model on.
dim (str, int) – Whether to build 2D or 3D network (if str, can be “2”, “2d”, “3D”, etc.)
- References:
- forward(x, sigma=None, **kwargs)[source]#
Run the denoiser on noisy image. The noise level is not used in this denoiser.
- Parameters:
x (torch.Tensor) – noisy image.
sigma (float) – noise level (not used).
Examples using UNet:#
Imaging inverse problems with adversarial networks
Self-supervised learning from incomplete measurements of multiple operators.
Self-supervised denoising with the Neighbor2Neighbor loss.
Self-supervised denoising with the Generalized R2R loss.
Self-supervised learning with measurement splitting
