neural_de.transformations.diffusion.unet package

Submodules

neural_de.transformations.diffusion.unet.attention_block module

class neural_de.transformations.diffusion.unet.attention_block.AttentionBlock(channels, num_heads=1, num_head_channels=-1, use_checkpoint=False, use_new_attention_order=False)

Bases: Module

An attention block that allows spatial positions to attend to each other.

Originally ported from here, but adapted to the N-d case. https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

neural_de.transformations.diffusion.unet.downsample module

class neural_de.transformations.diffusion.unet.downsample.Downsample(channels, use_conv, dims=2, out_channels=None)

Bases: Module

A downsampling layer with an optional convolution.

Parameters:

• channels – channels in the inputs and outputs.

• use_conv – a bool determining if a convolution is applied.

• dims – determines if the signal is 1D, 2D, or 3D. If 3D, then downsampling occurs in the inner-two dimensions.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

neural_de.transformations.diffusion.unet.nn module

class neural_de.transformations.diffusion.unet.nn.CheckpointFunction(*args, **kwargs)

Bases: Function

static backward(ctx, *output_grads)

Defines a formula for differentiating the operation with backward mode automatic differentiation (alias to the vjp function).

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computated w.r.t. the output.

static forward(ctx, run_function, length, *args)

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

• It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

• See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

• The forward no longer accepts a ctx argument.

• Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

• See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

class neural_de.transformations.diffusion.unet.nn.GroupNorm32(num_groups, num_channels, eps=1e-05, affine=True, device=None, dtype=None)

Bases: GroupNorm

affine: bool

eps: float

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

num_channels: int

num_groups: int

neural_de.transformations.diffusion.unet.nn.avg_pool_nd(dims, *args, **kwargs)

Create a 1D, 2D, or 3D average pooling module.

neural_de.transformations.diffusion.unet.nn.checkpoint(func, inputs, params, flag)

Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pass.

Parameters:

• func – the function to evaluate.

• inputs – the argument sequence to pass to func.

• params – a sequence of parameters func depends on but does not explicitly take as arguments.

• flag – if False, disable gradient checkpointing.

neural_de.transformations.diffusion.unet.nn.conv_nd(dims, *args, **kwargs)

Create a 1D, 2D, or 3D convolution module.

neural_de.transformations.diffusion.unet.nn.linear(*args, **kwargs)

Create a linear module.

neural_de.transformations.diffusion.unet.nn.normalization(channels)

Make a standard normalization layer.

Parameters:

channels – number of input channels.

Returns:

an nn.Module for normalization.

neural_de.transformations.diffusion.unet.nn.timestep_embedding(timesteps, dim, max_period=10000)

Create sinusoidal timestep embeddings.

Parameters:

• timesteps – a 1-D Tensor of N indices, one per batch element. These may be fractional.

• dim – the dimension of the output.

• max_period – controls the minimum frequency of the embeddings.

Returns:

an [N x dim] Tensor of positional embeddings.

neural_de.transformations.diffusion.unet.nn.zero_module(module)

Zero out the parameters of a module and return it.

neural_de.transformations.diffusion.unet.qkv_attention module

class neural_de.transformations.diffusion.unet.qkv_attention.QKVAttention(n_heads)

Bases: Module

A module which performs QKV attention and splits in a different order.

forward(qkv)

Apply QKV attention.

Parameters:

qkv – an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.

Returns:

an [N x (H * C) x T] tensor after attention.

training: bool

neural_de.transformations.diffusion.unet.qkv_attention_legacy module

class neural_de.transformations.diffusion.unet.qkv_attention_legacy.QKVAttentionLegacy(n_heads)

Bases: Module

A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping

forward(qkv)

Apply QKV attention.

Parameters:

qkv – an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.

Returns:

an [N x (H * C) x T] tensor after attention.

training: bool

neural_de.transformations.diffusion.unet.res_block module

class neural_de.transformations.diffusion.unet.res_block.ResBlock(channels, emb_channels, dropout, out_channels=None, use_conv=False, use_scale_shift_norm=False, dims=2, use_checkpoint=False, up=False, down=False)

Bases: TimestepBlock

A residual block that can optionally change the number of channels.

Parameters:

• channels – the number of input channels.

• emb_channels – the number of timestep embedding channels.

• dropout – the rate of dropout.

• out_channels – if specified, the number of out channels.

• use_conv – if True and out_channels is specified, use a spatial convolution instead of a smaller 1x1 convolution to change the channels in the skip connection.

• dims – determines if the signal is 1D, 2D, or 3D.

• use_checkpoint – if True, use gradient checkpointing on this module.

• up – if True, use this block for upsampling.

• down – if True, use this block for downsampling.

forward(x, emb)

Apply the block to a Tensor, conditioned on a timestep embedding.

Parameters:

• x – an [N x C x …] Tensor of features.

• emb – an [N x emb_channels] Tensor of timestep embeddings.

Returns:

an [N x C x …] Tensor of outputs.

training: bool

neural_de.transformations.diffusion.unet.timestep_block module

class neural_de.transformations.diffusion.unet.timestep_block.TimestepBlock(*args, **kwargs)

Bases: Module

Any module where forward() takes timestep embeddings as a second argument.

abstract forward(x, emb)

Apply the module to x given emb timestep embeddings.

training: bool

neural_de.transformations.diffusion.unet.timestep_embed_sequential module

class neural_de.transformations.diffusion.unet.timestep_embed_sequential.TimestepEmbedSequential(*args: Module)

class neural_de.transformations.diffusion.unet.timestep_embed_sequential.TimestepEmbedSequential(arg: OrderedDict[str, Module])

Bases: Sequential, TimestepBlock

A sequential module that passes timestep embeddings to the children that support it as an extra input.

forward(x, emb)

Apply the module to x given emb timestep embeddings.

neural_de.transformations.diffusion.unet.unet_model module

class neural_de.transformations.diffusion.unet.unet_model.UNetModel(in_channels, out_channels, config, logger=None)

Bases: Module

The full UNet model with attention and timestep embedding.

convert_to_fp16()

Convert the torso of the model to float16.

convert_to_fp32()

Convert the torso of the model to float32.

forward(x, timesteps, y=None)

Apply the model to an input batch.

Parameters:

• x – an [N x C x …] Tensor of inputs.

• timesteps – a 1-D batch of timesteps.

• y – an [N] Tensor of labels, if class-conditional.

Returns:

an [N x C x …] Tensor of outputs.

training: bool

neural_de.transformations.diffusion.unet.upsample module

class neural_de.transformations.diffusion.unet.upsample.Upsample(channels, use_conv, dims=2, out_channels=None)

Bases: Module

An upsampling layer with an optional convolution.

Parameters:

• channels – channels in the inputs and outputs.

• use_conv – a bool determining if a convolution is applied.

• dims – determines if the signal is 1D, 2D, or 3D. If 3D, then upsampling occurs in the inner-two dimensions.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

neural_de.transformations.diffusion.unet.utils module

neural_de.transformations.diffusion.unet.utils.convert_module_to_f16(ll)

Convert primitive modules to float16.

neural_de.transformations.diffusion.unet.utils.convert_module_to_f32(ll)

Convert primitive modules to float32, undoing convert_module_to_f16().

Module contents

This module contains the implementation of the unet architecture used par DiffPure.