data
UNetConfig
extends
ModelConfigUNetConfig(num_classes: int = 2, in_channels: int = 1, base_channels: int = 64, depth: int = 4, bilinear: bool = False, dropout: float = 0.0, dim: Literal[2, 3] = 2, block: Literal['basic', 'res'] = 'basic')Configuration for U-Net.
U-Net is a fully convolutional encoder-decoder architecture with skip connections between corresponding encoder and decoder stages. The architecture was originally proposed for biomedical image segmentation.
Channel schedule (base_channels=64, depth=4): Encoder: 64 → 128 → 256 → 512 Bottleneck: 1024 Decoder: 512 → 256 → 128 → 64
Parameters
num_classesint= 2Number of output segmentation classes.
in_channelsint= 1Number of input image channels.
base_channelsint= 64Feature channels at the first encoder stage.
Doubles at each depth level.
depthint= 4Number of encoder/decoder stages (excluding bottleneck).
bilinearbool= FalseIf True, use bilinear upsampling + Conv2d;
otherwise use ConvTranspose2d for learned upsampling.
dropoutfloat= 0.0Dropout probability applied in DoubleConv blocks.
Notes
Encoder: depth × (DoubleConv + MaxPool2d) — halves spatial resolution. Bottleneck: DoubleConv at the deepest level. Decoder: depth × (Upsample + skip-cat + DoubleConv). Head: Conv2d(base_channels, num_classes, 1).