UNet2DModel for NCSN

UNet2DModelForNCSN is a 2D U-Net model suitable for Noise Conditional Score Networks (NCSN) proposed by by Yang Song and Stefano Ermon in the paper Generative Modeling by Estimating Gradients of the Data Distribution. The model inherits from the diffuser’s UNet2DModel and is designed to be used with the AnnealedLangevinDynamicsScheduler scheduler.

The abstract of the paper is the following:

We introduce a new generative model where samples are produced via Langevin dynamics using gradients of the data distribution estimated with score matching. Because gradients can be ill-defined and hard to estimate when the data resides on low-dimensional manifolds, we perturb the data with different levels of Gaussian noise, and jointly estimate the corresponding scores, i.e., the vector fields of gradients of the perturbed data distribution for all noise levels. For sampling, we propose an annealed Langevin dynamics where we use gradients corresponding to gradually decreasing noise levels as the sampling process gets closer to the data manifold. Our framework allows flexible model architectures, requires no sampling during training or the use of adversarial methods, and provides a learning objective that can be used for principled model comparisons. Our models produce samples comparable to GANs on MNIST, CelebA and CIFAR-10 datasets, achieving a new state-of-the-art inception score of 8.87 on CIFAR-10. Additionally, we demonstrate that our models learn effective representations via image inpainting experiments.

UNet2DModelForNCSN

class ncsn.unet.UNet2DModelForNCSN(sigma_min, sigma_max, num_train_timesteps, sample_size=None, in_channels=3, out_channels=3, center_input_sample=False, time_embedding_type='positional', freq_shift=0, flip_sin_to_cos=True, down_block_types=('DownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D'), up_block_types=('AttnUpBlock2D', 'AttnUpBlock2D', 'AttnUpBlock2D', 'UpBlock2D'), block_out_channels=(224, 448, 672, 896), layers_per_block=2, mid_block_scale_factor=1, downsample_padding=1, downsample_type='conv', upsample_type='conv', dropout=0.0, act_fn='silu', attention_head_dim=8, norm_num_groups=32, attn_norm_num_groups=None, norm_eps=1e-05, resnet_time_scale_shift='default', add_attention=True, class_embed_type=None, num_class_embeds=None)[source]

A 2D UNet model for Noise Conditional Score Networks (NCSN).

This model inherits from UNet2DModel, which is a 2D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. This model also inherits from ModelMixin. Check the superclass documentation for it’s generic methods implemented for all models (such as downloading or saving).

Parameters:

  • sigma_min (float) – Minimum standard deviation for the isotropic Gaussian noise.

  • sigma_max (float) – Maximum standard deviation for the isotropic Gaussian noise.

  • sample_size (int or Tuple[int, int], optional, defaults to None) – Height and width of input/output sample. Dimensions must be a multiple of 2 ** (len(block_out_channels) - 1).

  • in_channels (int, optional, defaults to 3) – Number of channels in the input sample.

  • out_channels (int, optional, defaults to 3) – Number of channels in the output.

  • center_input_sample (bool, optional, defaults to False) – Whether to center the input sample.

  • time_embedding_type (str, optional, defaults to “positional”) – Type of time embedding to use.

  • freq_shift (int, optional, defaults to 0) – Frequency shift for Fourier time embedding.

  • flip_sin_to_cos (bool, optional, defaults to True) – Whether to flip sin to cos for Fourier time embedding.

  • down_block_types (Tuple[str], optional, defaults to (“DownBlock2D”, “AttnDownBlock2D”, “AttnDownBlock2D”, “AttnDownBlock2D”)) – Tuple of downsample block types.

  • mid_block_type (str, optional, defaults to “UNetMidBlock2D”) – Block type for middle of UNet, it can be either UNetMidBlock2D or UnCLIPUNetMidBlock2D.

  • up_block_types (Tuple[str], optional, defaults to (“AttnUpBlock2D”, “AttnUpBlock2D”, “AttnUpBlock2D”, “UpBlock2D”)) – Tuple of upsample block types.

  • block_out_channels (Tuple[int], optional, defaults to (224, 448, 672, 896)) – Tuple of block output channels.

  • layers_per_block (int, optional, defaults to 2) – The number of layers per block.

  • mid_block_scale_factor (float, optional, defaults to 1) – The scale factor for the mid block.

  • downsample_padding (int, optional, defaults to 1) – The padding for the downsample convolution.

  • downsample_type (str, optional, defaults to conv) – The downsample type for downsampling layers. Choose between “conv” and “resnet”

  • upsample_type (str, optional, defaults to conv) – The upsample type for upsampling layers. Choose between “conv” and “resnet”

  • dropout (float, optional, defaults to 0.0) – The dropout probability to use.

  • act_fn (str, optional, defaults to “silu”) – The activation function to use.

  • attention_head_dim (int, optional, defaults to 8) – The attention head dimension.

  • norm_num_groups (int, optional, defaults to 32) – The number of groups for normalization.

  • attn_norm_num_groups (int, optional, defaults to None) – If set to an integer, a group norm layer will be created in the mid block’s [Attention] layer with the given number of groups. If left as None, the group norm layer will only be created if resnet_time_scale_shift is set to default, and if created will have norm_num_groups groups.

  • norm_eps (float, optional, defaults to 1e-5) – The epsilon for normalization.

  • resnet_time_scale_shift (str, optional, defaults to “default”) – Time scale shift config for ResNet blocks (see ResnetBlock2D). Choose from default or scale_shift.

  • class_embed_type (str, optional, defaults to None) – The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, “timestep”, or “identity”.

  • num_class_embeds (int, optional, defaults to None) – Input dimension of the learnable embedding matrix to be projected to time_embed_dim when performing class conditioning with class_embed_type equal to None.

__init__(sigma_min, sigma_max, num_train_timesteps, sample_size=None, in_channels=3, out_channels=3, center_input_sample=False, time_embedding_type='positional', freq_shift=0, flip_sin_to_cos=True, down_block_types=('DownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D'), up_block_types=('AttnUpBlock2D', 'AttnUpBlock2D', 'AttnUpBlock2D', 'UpBlock2D'), block_out_channels=(224, 448, 672, 896), layers_per_block=2, mid_block_scale_factor=1, downsample_padding=1, downsample_type='conv', upsample_type='conv', dropout=0.0, act_fn='silu', attention_head_dim=8, norm_num_groups=32, attn_norm_num_groups=None, norm_eps=1e-05, resnet_time_scale_shift='default', add_attention=True, class_embed_type=None, num_class_embeds=None)[source]

Initialize internal Module state, shared by both nn.Module and ScriptModule.