Skip to content

UNet Config

Source

UNet Pydantic model.

UNetConfig

Bases: ArchitectureConfig

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

Attributes:

Name Type Description
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.

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:

Type Description
bool

Whether the model is 3D or not.

model_dump(**kwargs)

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

Parameters:

Name Type Description Default
**kwargs Any

Additional keyword arguments from Pydantic BaseModel model_dump method.

{}

Returns:

Type Description
{str: Any}

Model as a dictionary.

set_3D(is_3D)

Set 3D model by setting the conv_dims parameters.

Parameters:

Name Type Description Default
is_3D bool

Whether the algorithm is 3D or not.

required

validate_num_channels_init(num_channels_init) classmethod

Validate that num_channels_init is even.

Parameters:

Name Type Description Default
num_channels_init int

Number of channels.

required

Returns:

Type Description
int

Validated number of channels.

Raises:

Type Description
ValueError

If the number of channels is odd.