zetta_utils.convnet

zetta_utils.convnet.architecture

class zetta_utils.convnet.architecture.ConvBlock(num_channels, kernel_sizes, conv=<class 'torch.nn.modules.conv.Conv2d'>, strides=None, activation=<class 'torch.nn.modules.activation.LeakyReLU'>, normalization=None, skips=None, normalize_last=False, activate_last=False, paddings='same', padding_modes='zeros', activation_mode='post')[source]

A block with sequential convolutions with activations, optional normalization and residual connections.

Parameters

  • num_channels (Sequence[int]) – list of integers specifying the number of channels of each convolution. For example, specification [1, 2, 3] will correspond to a sequence of 2 convolutions, where the first one has 1 input channel and 2 output channels and the second one has 2 input channels and 3 output channels.

  • conv (Callable[…, _ConvNd]) – Constructor for convolution layers.

  • activation (Callable[[], Module]) – Constructor for activation layers.

  • normalization (Optional[Callable[[int], Module]]) – Constructor for normalization layers. Normalization will be applied after convolution before activation.

  • kernel_sizes (Union[Sequence[int], Sequence[Sequence[int]]]) – Convolution kernel sizes. When specified as a sequence of integerss it will be passed as k parameter to all convolution constructors. When specified as a sequence of sequence of integers, each item in the outer sequence will be passed as k to the corresponding convolution in order, and the outer sequence length must match the number of convolutions.

  • strides (Union[Sequence[int], Sequence[Sequence[int]], None]) – Convolution strides. When specified as a sequence of integers it will be passed as the stride parameter to all convolution constructors. When specified as a sequence of sequences, each item in the outer sequence will be passed as stride to the corresponding convolution in order, and the outer sequence length must match the number of convolutions.

  • paddings (Union[Literal[‘same’, ‘valid’], Sequence[int], Sequence[Union[Literal[‘same’, ‘valid’], Sequence[int]]]]) – Convolution padding sizes. When specified as a single string, or a sequence of integers, it will be passed as the padding parameter to all convolution constructors. When specified as a sequence of strings or sequence of int sequences, each item in the outer sequence will be passed as padding to the corresponding convolution in order, and the outer sequence length must match the number of convolutions.

  • skips (UnionType[dict[str, int], None]) – Specification for residual skip connection. For example, skips={"1": 3} specifies a single residual skip connection from the output of the first convolution (index 1) to the output of third convolution (index 3). 0 specifies the input to the first layer.

  • normalize_last (bool) – Whether to apply normalization after the last layer.

  • activate_last (bool) – Whether to apply activation after the last layer.

  • padding_modes (Union[Literal[‘zeros’, ‘reflect’, ‘replicate’, ‘circular’], Sequence[Literal[‘zeros’, ‘reflect’, ‘replicate’, ‘circular’]]]) – Convolution padding modes. When specified as a single string, it will be passed as the padding mode parameter to all convolution constructors. When specified as a sequence of strings, each item in the sequence will be passed as padding mode to the corresponding convolution in order. The list length must match the number of convolutions.

  • activation_mode (Literal[‘pre’, ‘post’]) – Whether to use pre-activation (norm. -> act. -> conv.) or post-activation (conv. -> norm. -> act.) for residual skip connections (He et al. 2016).

class zetta_utils.convnet.architecture.UNet(list_num_channels, kernel_sizes, conv=<class 'torch.nn.modules.conv.Conv2d'>, strides=None, downsample=functools.partial(<class 'torch.nn.modules.pooling.AvgPool2d'>, kernel_size=2), upsample=functools.partial(<class 'torch.nn.modules.upsampling.Upsample'>, scale_factor=2), activation=<class 'torch.nn.modules.activation.LeakyReLU'>, normalization=None, paddings='same', skips=None, normalize_last=False, activate_last=False, padding_modes='zeros', unet_skip_mode='sum', activation_mode='post')[source]

A basic UNet that uses ConvBlocks. Note that the downsamples are applied after the regular convolutions, while the upsamples are applied before the regular convolutions.

Parameters

  • list_num_channels (Sequence[Sequence[int]]) – List of lists of integers specifying the number of channels of each convolution. There must be an odd number of lists, for going up and going down, and the middle. Each list is used to generate the ConvBlocks.

  • conv (Callable[…, _ConvNd]) – Constructor for convolution layers.

  • downsample (Callable[…, Module]) – Constructor for downsample layers. Must include strides and kernel sizes. Activation and normalization are included by default and cannot be turned off.

  • upsample (Callable[…, Module]) – Constructor for upsample layers. Must include strides and kernel sizes. Activation and normalization are included by default and cannot be turned off.

  • activation (Callable[[], Module]) – Constructor for activation layers.

  • normalization (Optional[Callable[[int], Module]]) – Constructor for normalization layers. Normalization will be applied after convolution before activation.

  • kernel_sizes (Union[Sequence[int], Sequence[Sequence[int]]]) – List of convolution kernel sizes. Will be passed directly to zetta_utils.convnet.architecture.ConvBlock constructors.

  • strides (Union[Sequence[int], Sequence[Sequence[int]], None]) – List of convolution strides. Will be passed directly to zetta_utils.convnet.architecture.ConvBlock constructors.

  • paddings (Union[Literal[‘same’, ‘valid’], Sequence[int], Sequence[Union[Literal[‘same’, ‘valid’], Sequence[int]]]]) – List of convolution padding sizes. Will be passed directly to zetta_utils.convnet.architecture.ConvBlock constructors.

  • skips (UnionType[dict[str, int], None]) – Specifications for residual skip connections as documented in ConvBlock. Will be passed directly to zetta_utils.convnet.architecture.ConvBlock constructors.

  • normalize_last (bool) – Whether to apply normalization after the last layer. Note that normalization is included by default within and at the end of the convblocks when the normalization layer is specified, and cannot be turned off.

  • activate_last (bool) – Whether to apply activation after the last layer. Note that activation is included by default within and at the end of the convblocks and cannot be turned off.

  • padding_modes (Union[Literal[‘zeros’, ‘reflect’, ‘replicate’, ‘circular’], Sequence[Literal[‘zeros’, ‘reflect’, ‘replicate’, ‘circular’]]]) – Will be passed directly to zetta_utils.convnet.architecture.ConvBlock constructors.

  • unet_skip_mode (Literal[‘sum’, ‘concat’]) – Whether skip connections are realized as residual connections (sum, default) or concatenated skip connections (concat).