UnetEncoder

class UnetEncoder(num_input_channels, layer_output_channels)[source]

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)

Methods

forward

Logic for using layers defined in init.

Attributes

forward(x)[source]

Logic for using layers defined in init.

This method defines how layers are used in forward operation.

Parameters

  • x (torch.Tensor) – Input images- the tensor is of the shape

  • NCHW. –

Returns

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

Return type

list