U-Net

A typed encoder-decoder with one downsampling stage, one upsampling stage, and a channel-wise skip connection. The spatial rank is a parameter: the same constructor applies to sequences, images, volumes, and higher-dimensional grids.

The reusable operations remain visible in the type. Pooling determines the bottleneck extent, transpose convolution returns it to the input extent, and branch concatenation changes the channel count from baseChannels to baseChannels + baseChannels.

Reference: O. Ronneberger, P. Fischer, and T. Brox, "U-Net: Convolutional Networks for Biomedical Image Segmentation," MICCAI 2015.

structure TorchLean.nn.models.UNet.Config (d : ℕ) : Type

Configuration for a one-level U-Net over d spatial axes.

• inputChannels : ℕ

  Number of channels in each input sample.

• baseChannels : ℕ

  Width of the full-resolution feature map.

• outputChannels : ℕ

  Number of channels produced by the output projection.

• spatial : Tensor ℕ [d]

  Size of each input axis.

• pooling : Pooling.Config d

  Downsampling operation.

• kernelRadius : Tensor ℕ [d]

  Radius of the same-padding kernels used at both resolutions.

  A radius of 1 gives kernel size 3; these convolutions preserve their grids automatically.

• upsampling : Convolution.Geometry d

  Geometry used by the transpose convolution that restores full resolution.

def TorchLean.nn.models.UNet.Config.pooledSpatial {d : ℕ} (config : Config d) : Tensor ℕ [d]

Grid size after the downsampling stage.

def TorchLean.nn.models.UNet.Config.validate {d : ℕ} (config : Config d) : Except String Unit

Validate every channel width and downsample/upsample geometry before construction.

@[reducible, inline]

abbrev TorchLean.nn.models.UNet.Config.inputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Input shape with arbitrary batch axes.

@[reducible, inline]

abbrev TorchLean.nn.models.UNet.Config.outputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Output shape with the same batch axes as the input.

def TorchLean.nn.models.unet {d : ℕ} (config : UNet.Config d) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build a one-level U-Net over an arbitrary spatial rank.

Same-padding convolutions preserve both resolutions automatically. Model validation rejects an empty input or pooled grid, or a pooling/upsampling pair that does not restore the input grid.