OutHeadConv2D

class aixd.mlmodel.architecture.heads.OutHeadConv2D(*args: Any, **kwargs: Any)[source]

Bases: OutHeadConv

Convolutional Network for decoding spatially strucured (image-like) data in 2 dimensions.

Parameters:
  • in_channels (int) – The number of input channels.

  • target_shape (Tuple[int]) – Final shape that the output data should have.

  • latent_dims (List[int]) – A list of integers representing the dimensions of the latent space. The length of the sequence defines the number of layers.

  • activation (Union[nn.Module, str]) – The activation function to use.

  • out_activation (Union[nn.Module, str]) – The activation function to use in the last layer.

  • attn_block_indices (List[int], optional, default=[]) – Sequence of indices where to insert an attention layer in the encoding blocks.

Inherited Methods