modelzoo.vision.pytorch.unet.layers.UNetBlock.UNetBlock#

class modelzoo.vision.pytorch.unet.layers.UNetBlock.UNetBlock[source]#

Bases: torch.nn.Module

Parameters
  • in_channels (int) – The input channel dimension before passing through the block.

  • out_channels (int) – The output channel dimension after passing through the block.

  • encoder (bool) – Flag if the block is part of the encoder section of the model. Returns the block-output if True, and a tuple of block-output, block-output before downsampling if False.

  • convs_per_block ([str]) – List of conv specifications for each conv in the block.

  • bias (bool) – Flag to use bias vectors.

  • norm_layer (nn.Module) – Desired normalization layer.

  • norm_kwargs (dict) – A dictionary of the arguments to pass to the constructor of the normalization layer.

  • act (str) – Activation to use.

  • skip_connect (bool) – Flag for if the model concatenates encoder outputs to decoder inputs.

  • residual_blocks (bool) – Flag for using residual connections.

  • downscale_method (str) – Downscaling method at the end of the block.

Methods

forward

__call__(*args: Any, **kwargs: Any) Any#

Call self as a function.

__init__(in_channels, out_channels, encoder, convs_per_block, bias, norm_layer, norm_kwargs, act='relu', skip_connect=True, residual_blocks=False, downscale_method='max_pool', use_conv3d=False, downscale_first_conv=False, downscale=True)[source]#
static __new__(cls, *args: Any, **kwargs: Any) Any#