unet

Implementation of a flexible U-net.

Originally inspired by: https://github.com/milesial/Pytorch-UNet

Classes:

• ConvBlock –

  Convolution block: conv => BN => act.

• DoubleConv –

  Double convolution (conv => BN => ReLU) * 2.

• DownBlock –

  Down-scaling block.

• UNet –

  U-net model.

• UpBlock –

  Up-scaling block.

ConvBlock

ConvBlock(
    in_ch: int,
    out_ch: int,
    kernel_size: int,
    n_dims: int = 2,
    stride: int = 1,
    dilation: int = 1,
    pad_mode: str = "replicate",
    residual: bool = False,
    bias: bool = True,
    last_block: bool = False,
)

Bases: Sequential

Convolution block: conv => BN => act.

Source code in src/autoden/models/unet.py

def __init__(
    self,
    in_ch: int,
    out_ch: int,
    kernel_size: int,
    n_dims: int = 2,
    stride: int = 1,
    dilation: int = 1,
    pad_mode: str = "replicate",
    residual: bool = False,
    bias: bool = True,
    last_block: bool = False,
):
    pad_size = (kernel_size - 1) // 2 + (dilation - 1)
    if last_block:
        post_conv = []
    else:
        post_conv = [NDBatchNorm[n_dims](out_ch), nn.LeakyReLU(0.2, inplace=True)]
    super().__init__(
        NDConv[n_dims](
            in_ch,
            out_ch,
            kernel_size=kernel_size,
            stride=stride,
            dilation=dilation,
            padding=pad_size,
            padding_mode=pad_mode,
            bias=bias,
        ),
        *post_conv,
    )
    if residual and in_ch != out_ch:
        print(f"Warning: Residual connections not available when {in_ch=} is different from {out_ch=}")
        residual = False
    self.residual = residual

DoubleConv

DoubleConv(
    in_ch: int,
    out_ch: int,
    n_dims: int = 2,
    pad_mode: str = "replicate",
)

Bases: Sequential

Double convolution (conv => BN => ReLU) * 2.

Source code in src/autoden/models/unet.py

def __init__(self, in_ch: int, out_ch: int, n_dims: int = 2, pad_mode: str = "replicate"):
    super().__init__(
        ConvBlock(in_ch, out_ch, n_dims=n_dims, kernel_size=3, pad_mode=pad_mode),
        ConvBlock(out_ch, out_ch, n_dims=n_dims, kernel_size=1),
    )

DownBlock

DownBlock(
    in_ch: int,
    out_ch: int,
    n_dims: int = 2,
    bilinear: bool = True,
    pad_mode: str = "replicate",
)

Bases: Sequential

Down-scaling block.

Source code in src/autoden/models/unet.py

def __init__(self, in_ch: int, out_ch: int, n_dims: int = 2, bilinear: bool = True, pad_mode: str = "replicate"):
    if bilinear:
        down_block = [NDPool[n_dims](2)]
    else:
        down_block = [ConvBlock(in_ch, in_ch, n_dims=n_dims, kernel_size=2, stride=2)]
    super().__init__(
        *down_block,
        DoubleConv(in_ch, out_ch, n_dims=n_dims, pad_mode=pad_mode),
    )
    self.n_dims = n_dims
    self.pad_mode = pad_mode.lower()

UNet

UNet(
    n_channels_in: int,
    n_channels_out: int,
    n_features: int = 32,
    n_levels: int = 3,
    n_dims: int = 2,
    n_channels_skip: int | None = None,
    bilinear: bool = True,
    pad_mode: str = "replicate",
    device: str = "cuda" if is_available() else "cpu",
    verbose: bool = False,
)

Bases: Module

U-net model.

Source code in src/autoden/models/unet.py

def __init__(
    self,
    n_channels_in: int,
    n_channels_out: int,
    n_features: int = 32,
    n_levels: int = 3,
    n_dims: int = 2,
    n_channels_skip: int | None = None,
    bilinear: bool = True,
    pad_mode: str = "replicate",
    device: str = "cuda" if pt.cuda.is_available() else "cpu",
    verbose: bool = False,
):
    init_params = locals()
    del init_params["self"]
    del init_params["__class__"]

    super().__init__()
    self.init_params = init_params
    self.device = device

    if pad_mode.lower() not in PAD_MODES:
        raise ValueError(f"Padding mode {pad_mode} should be one of {PAD_MODES}")

    encoder, decoder = _compute_architecture(
        n_levels=n_levels, n_features=n_features, n_skip=n_channels_skip, verbose=verbose
    )

    self.in_layer = DoubleConv(n_channels_in, n_features, n_dims=n_dims, pad_mode=pad_mode)
    self.encoder_layers = nn.ModuleList(
        [DownBlock(*lvl, n_dims=n_dims, bilinear=bilinear, pad_mode=pad_mode) for lvl in encoder]
    )
    self.decoder_layers = nn.ModuleList(
        [UpBlock(*lvl, n_dims=n_dims, linear=bilinear, pad_mode=pad_mode) for lvl in decoder]
    )
    self.out_layer = ConvBlock(
        n_features, n_channels_out, n_dims=n_dims, kernel_size=1, pad_mode=pad_mode, last_block=True
    )

    self.to(self.device)

UpBlock

UpBlock(
    in_ch: int,
    skip_ch: int | None,
    out_ch: int,
    n_dims: int = 2,
    linear: bool = True,
    pad_mode: str = "replicate",
)

Bases: Module

Up-scaling block.

Source code in src/autoden/models/unet.py

def __init__(
    self, in_ch: int, skip_ch: int | None, out_ch: int, n_dims: int = 2, linear: bool = True, pad_mode: str = "replicate"
):
    super().__init__()
    self.skip_ch = skip_ch
    self.n_dims = n_dims

    # Bilinear up-sampling tends to give better results, and use fewer weights
    if linear:
        self.up_block = nn.Upsample(scale_factor=2, mode=NDUpsampling[n_dims], align_corners=True)
    else:
        self.up_block = NDConvT[n_dims](in_ch, in_ch, kernel_size=2, stride=2)

    if skip_ch is not None:
        n_skip = skip_ch
        if skip_ch > 0:
            self.skip_block = ConvBlock(in_ch, skip_ch, n_dims=n_dims, kernel_size=1, pad_mode=pad_mode)
    else:
        n_skip = in_ch
    self.conv_block = DoubleConv(in_ch + n_skip, out_ch, n_dims=n_dims, pad_mode=pad_mode)