unet blocks

class src.blocks.ConvBlock(*args: Any, **kwargs: Any)

U-Net Convolution block

Parameters

in_channels, out_channels: int

Number of input and desired output channels.

kernel_size: `tuple` or int, optional

Kernel size to use for the convolution. Default is (3, 3).

stride: `tuple` or int, optional

Stride to use for the convolution. Default is 1.

depth: int, optional

Amount of convolutions to perform. Default is 2.

activation: str, optional

Activation to use after every convolution. If ‘relu’ is not specified, no activation will be used. Default is ‘relu’.

batch_norm: bool, optional

Whether or not to use batch normalization after the activation. Default is True.

padding: int or bool, optional

Padding to apply before convoluting. Default is 0.

Attributes

net: nn.Sequential

Convolution block.

Methods

__call__(*args, **kwargs)

Call self as a function.

forward

class src.blocks.DeconvBlock(*args: Any, **kwargs: Any)

U-Net Deconvolution block

Parameters

in_channels, out_channels: int

Number of input and desired output channels.

up: str, optional

Mode of the upconvolution. Can be either ‘upsample’, applying upsampling to increase image dimensions, or ‘upconv’, applying a transposed convolution (nn.ConvTranspose2d). Default is ‘upsample’.

scale_factor: `tuple` or int, optional

Scale of the upscaled image compared to the input image. When applying transposed convolution, this will also be used as stride. Default is (2, 2).

mode: str, optional

Upsampling algorithm. This only applies when using upsample mode. For 4D tensors this may only be bilinear. Default is ‘bilinear’.

align_corners: bool, optional

Whether or not to align image corners when upsampling. This only applies when using upsampling mode. Default is False.

**kwargs

Key-word arguments to be passed to the ConvBlock constructor when creating the convolution block.

Attributes

up: nn.Sequential

Upscaling block, using either upsampling or transposed convolution

conv: ConvBlock

Convolution block

Methods

__call__(*args, **kwargs)	Call self as a function.
center_crop(batch, target_size)	Center-crop image to specified size

forward