modelzoo.vision.pytorch.dit.layers.vae.ResNetBlock2D.ResnetBlock2D#

class modelzoo.vision.pytorch.dit.layers.vae.ResNetBlock2D.ResnetBlock2D[source]#

Bases: torch.nn.Module

A Resnet block.

Parameters

  • in_channels (int) – The number of channels in the input.

  • out_channels (int, optional, default to be None) – The number of output channels for the first conv2d layer. If None, same as in_channels.

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

  • temb_channels (int, optional, default to 512) – the number of channels in timestep embedding.

  • groups (int, optional, default to 32) – The number of groups to use for the first normalization layer.

  • groups_out (int, optional, default to None) – The number of groups to use for the second normalization layer. if set to None, same as groups.

  • eps (float, optional, defaults to 1e-6) – The epsilon to use for the normalization.

  • non_linearity (str, optional, default to “swish”) – the activation function to use.

  • time_embedding_norm (str, optional, default to “default”) – Time scale shift config. By default, apply timestep embedding conditioning with a simple shift mechanism. Choose “scale_shift” or “ada_group” for a stronger conditioning with scale and shift.

  • kernal (torch.FloatTensor, optional, default to None) – FIR filter, see [~models.resnet.FirUpsample2D] and [~models.resnet.FirDownsample2D].

  • output_scale_factor (float, optional, default to be 1.0) – the scale factor to use for the output.

  • use_in_shortcut (bool, optional, default to True) – If True, add a 1x1 nn.conv2d layer for skip-connection.

  • up (bool, optional, default to False) – If True, add an upsample layer.

  • down (bool, optional, default to False) – If True, add a downsample layer.

  • conv_shortcut_bias (bool, optional, default to True) – If True, adds a learnable bias to the conv_shortcut output.

  • conv_2d_out_channels (int, optional, default to None) – the number of channels in the output. If None, same as out_channels.

Methods

forward

__call__(*args: Any, **kwargs: Any) → Any#

Call self as a function.

__init__(*, in_channels, out_channels=None, conv_shortcut=False, dropout=0.0, temb_channels=512, groups=32, groups_out=None, pre_norm=True, eps=1e-06, non_linearity='swish', time_embedding_norm='default', kernel=None, output_scale_factor=1.0, use_in_shortcut=None, up=False, down=False, conv_shortcut_bias: bool = True, conv_2d_out_channels: Optional[int] = None)[source]#
static __new__(cls, *args: Any, **kwargs: Any) → Any#

previous

modelzoo.vision.pytorch.dit.layers.vae.ResNetBlock2D.Downsample2D

next

modelzoo.vision.pytorch.dit.layers.vae.ResNetBlock2D.Upsample2D