unet(pretrained: bool = False, overrides: object = {})U-Net (Ronneberger et al., MICCAI 2015).
Builds the canonical 2-D U-Net: 4-level encoder / decoder with
base_channels = 64 (channel schedule 64 -> 128 -> 256 -> 512 ->
1024), in_channels = 1 (single-channel biomedical default), and
num_classes = 2. ConvTranspose2d upsampling in the decoder
(set bilinear=True to use bilinear interpolation instead).
Approximately 31M parameters.
Model Size
Parameters
pretrainedbool= FalseReserved for future pretrained-weight loading. Currently ignored.
**overridesobject= {}Keyword overrides forwarded into
UNetConfig —
num_classes, in_channels (e.g. 3 for RGB), base_channels,
depth, bilinear, dropout.Returns
UNetForSemanticSegmentationSegmentation model with the U-Net configuration applied (or
with overrides merged on top of it).
Notes
See Ronneberger et al., "U-Net: Convolutional Networks for Biomedical Image Segmentation", MICCAI 2015 (arXiv:1505.04597). The skip connections that copy the encoder feature maps onto the decoder are the architectural key idea.
Examples
>>> import lucid
>>> from lucid.models.vision.unet import unet
>>> model = unet(num_classes=4, in_channels=3)
>>> x = lucid.randn(1, 3, 256, 256)
>>> out = model(x)
>>> out.logits.shape
(1, 4, 256, 256)