Skip to content

UNet Config

Source

UNet Pydantic model.

UNetConfig

Bases: ArchitectureConfig

Pydantic model for a N2V(2)-compatible UNet.

Attributes:

  • depth (int) –

    Depth of the model, between 1 and 10 (default 2).

  • num_channels_init (int) –

    Number of filters of the first level of the network, should be even and minimum 8 (default 96).

architecture instance-attribute

Name of the architecture.

conv_dims = Field(default=2, validate_default=True) class-attribute instance-attribute

Dimensions (2D or 3D) of the convolutional layers.

depth = Field(default=2, ge=1, le=10, validate_default=True) class-attribute instance-attribute

Number of levels in the UNet.

final_activation = Field(default='None', validate_default=True) class-attribute instance-attribute

Final activation function.

in_channels = Field(default=1, ge=1, validate_default=True) class-attribute instance-attribute

Number of channels in the input to the model.

independent_channels = Field(default=True, validate_default=True) class-attribute instance-attribute

Whether information is processed independently in each channel, used to train channels independently.

n2v2 = Field(default=False, validate_default=True) class-attribute instance-attribute

Whether to use N2V2 architecture modifications, with blur pool layers and fewer skip connections.

num_channels_init = Field(default=32, ge=8, le=1024, validate_default=True) class-attribute instance-attribute

Number of convolutional filters in the first layer of the UNet.

num_classes = Field(default=1, ge=1, validate_default=True) class-attribute instance-attribute

Number of classes or channels in the model output.

residual = Field(default=False, validate_default=True) class-attribute instance-attribute

Whether to add a residual connection from the input to the output.

use_batch_norm = Field(default=True, validate_default=True) class-attribute instance-attribute

Whether to use batch normalization in the model.

get_num_input_channels()

Get the number of input channels.

Returns:

  • int –

    Number of input channels.

get_num_output_channels()

Get the number of output channels.

Returns:

  • int –

    Number of output channels.

is_3D()

Return whether the model is 3D or not.

This method is used in the NG configuration validation to check that the model dimensions match the data dimensions.

Returns:

  • bool –

    Whether the model is 3D or not.

model_dump(**kwargs)

Dump the model as a dictionary, ignoring the architecture keyword.

Parameters:

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments from Pydantic BaseModel model_dump method.

Returns:

  • {str: Any} –

    Model as a dictionary.

set_3D(is_3D)

Set 3D model by setting the conv_dims parameters.

Parameters:

  • is_3D (bool) –

    Whether the algorithm is 3D or not.

uses_batch_norm()

Return whether the model uses batch normalization.

Returns:

  • bool –

    Whether the model uses batch normalization.

validate_num_channels_init(num_channels_init) classmethod

Validate that num_channels_init is even.

Parameters:

  • num_channels_init (int) –

    Number of channels.

Returns:

  • int –

    Validated number of channels.

Raises:

  • ValueError –

    If the number of channels is odd.