UnetEncoder

tiatoolbox . models . architecture . unet . UnetEncoder

class UnetEncoder(num_input_channels, layer_output_channels)

Construct a basic UNet encoder.

This class builds a basic UNet encoder with batch normalization. The number of channels in each down-sampling block and the number of down-sampling levels are customisable.

Parameters:

• num_input_channels (int) – Number of channels in the input images.

• layer_output_channels (list) – A list of integers where each number defines the number of output channels at each down-sampling level.

Returns:

A pytorch model.

Return type:

model (torch.nn.Module)

Initialize UnetEncoder.

Methods

forward

Logic for using layers defined in init.

Attributes

training

forward(input_tensor)

Logic for using layers defined in init.

This method defines how layers are used in forward operation.

Parameters:

• input_tensor (torch.Tensor) – Input images, the tensor is in the shape of NCHW.

• self (UnetEncoder)

Returns:

A list of features for each down-sample block. Each feature tensor is of the shape NCHW.

Return type:

list