ADMUNet

class deepinv.models.ADMUNet(img_resolution=64, in_channels=3, out_channels=3, label_dim=0, augment_dim=0, model_channels=192, channel_mult=[1, 2, 3, 4], channel_mult_emb=4, num_blocks=3, attn_resolutions=[32, 16, 8], dropout=0.10, label_dropout=0, pretrained=None, device=None, *args, **kwargs)

Bases: Denoiser

Re-implementation of the architecture from the paper: Diffusion Models Beat GANs on Image Synthesis.

Equivalent to the original implementation by Dhariwal and Nichol, available at: openai/guided-diffusion. The architecture consists of a series of convolution layer, down-sampling residual blocks and up-sampling residual blocks with skip-connections. Each residual block has a self-attention mechanism with 64 channels per attention head with the up/down-sampling from BigGAN.. The noise level is embedded using Positional Embedding with optional augmentation linear layer.

Parameters:

• img_resolution (int) – Image spatial resolution at input/output.

• in_channels (int) – Number of color channels at input.

• out_channels (int) – Number of color channels at output.

• label_dim (int) – Number of class labels, 0 = unconditional.

• augment_dim (int) – Augmentation label dimensionality, 0 = no augmentation.

• model_channels (int) – Base multiplier for the number of channels.

• channel_mult (list[int]) – Per-resolution multipliers for the number of channels.

• channel_mult_emb (int) – Multiplier for the dimensionality of the embedding vector.

• num_blocks (int) – Number of residual blocks per resolution.

• attn_resolutions (list[int]) – List of resolutions with self-attention.

• dropout (float) – dropout probability used in residual blocks.

• label_dropout (float) – Dropout probability of class labels for classifier-free guidance.

• pretrained (str, None) – 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 (the default model is a conditional model trained on ImageNet at 64x64 resolution (imagenet64-cond) with default architecture). Finally, pretrained can also be set as a path to the user’s own pretrained weights. See pretrained-weights for more details.

• device (torch.device) – Instruct our module to be either on cpu or on gpu. Default to None, which suggests working on cpu.

forward(x, sigma, class_labels=None, augment_labels=None)

Run the denoiser on noisy image.

Parameters:

• x (torch.Tensor) – noisy image

• sigma (Union[torch.Tensor, float]) – noise level

• class_labels (torch.Tensor) – class labels

• augment_labels (torch.Tensor) – augmentation labels

Returns:

(torch.Tensor) denoised image.