Dynamic Unet that can use any pretrained model as a backbone.

Dynamic U-Net

This module builds a dynamic U-Net from any backbone pretrained on ImageNet, automatically inferring the intermediate sizes.

U-Net architecure

This is the original U-Net. The difference here is that the left part is a pretrained model.

class DynamicUnet[source]

DynamicUnet(encoder:Module, n_classes:int) :: Sequential

Builds a U-Net from a given encoder (that can be a pretrained model) and with a final output of n_classes. During the initialization, it uses Hooks to determine the intermediate features sizes by passing a dummy input throught the model.

class UnetBlock[source]

UnetBlock(up_in_c:int, x_in_c:int, hook:Hook) :: Module

Builds a U-Net block that receives the output of the last block to be upsampled (size up_in_c) and the activations features from an intermediate layer of the encoder (size x_in_c, this is the lateral connection). The hook is set to this intermediate layer to store the output needed for this block.