# unet

Python API: `luxonis_train.nodes.blocks.unet`

The encoder and decoder blocks of a U-Net, with and without skip connections.

## Classes

### BaseDecoder

Base class of a decoder that mirrors
[SimpleEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).

The decoder adds `1` in front of the encoder multipliers and walks the list backwards. Each pair of neighbour entries gives one
decoder step. The step doubles the size and maps the channels of the first entry to the channels of the second. For `[1, 2]`, the
steps map `2 * base_width` channels to `base_width`, then `base_width` to `base_width`. A `3x3` convolution then maps `base_width`
channels to `out_channels`.

Every step uses `"simple_upsample"` with `align_corners=True`, and a batch norm in its
[UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
The class defines no `forward`.
[SimpleDecoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
and
[UNetDecoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
add it.

#### Methods

##### init

```python
def __init__(base_width: int, out_channels: int, encoder_width_multipliers: list[float], n_convolutions: int, block: type[SimpleDecoderBlock | UNetDecoderBlock]):
```

Build the decoder steps and the output convolution.

Parameters

 * `base_width` (`int`): The `base_hidden_channels` of the encoder.
 * `out_channels` (`int`): The number of output channels.
 * `encoder_width_multipliers` (`list[float]`): The `width_multipliers` of the encoder.
 * `n_convolutions` (`int`): The number of `3x3`
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   layers in the stack of each step.
 * `block` (`type[SimpleDecoderBlock | UNetDecoderBlock]`): The class of the decoder steps.

#### Attributes

##### blocks

The decoder steps, from the smallest size to the largest.

##### final_conv

The `3x3` output convolution.

### BaseDecoderBlock

Base class of a decoder step: an
[UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
and a
[ConvStack](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md).

The class defines no `forward`.
[SimpleDecoderBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
and
[UNetDecoderBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
add it.

#### Methods

##### init

```python
def __init__(in_channels: int, out_channels: int, conv_in_channels: int, kernel_size: int, use_norm: bool, align_corners: bool, upsample_mode: Literal['simple_upsample', 'conv_upsample', 'conv_transpose'], n_repeats: int):
```

Build the upsampling block and the convolution stack.

Parameters

 * `in_channels` (`int`): The number of input channels.
 * `out_channels` (`int`): The number of output channels of the upsampling block and of the stack.
 * `conv_in_channels` (`int`): The number of input channels of the stack.
 * `kernel_size` (`int`): The kernel size of
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
   Only `"conv_transpose"` uses it.
 * `use_norm` (`bool`): Whether the `3x3`
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   of
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
   has a batch norm. The stack always has batch norms.
 * `align_corners` (`bool`): The `align_corners` option of the interpolation in
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
 * `upsample_mode` (`Literal['simple_upsample', 'conv_upsample', 'conv_transpose']`): The upsampling method of
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
   The factor is `2`.
 * `n_repeats` (`int`): The number of `3x3`
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   layers in the stack.

#### Attributes

##### conv

The convolution stack after the upsampling.

##### up

The block that upsamples the input.

### EncoderBlock

Encoder step of an optional `2x2` max pool and a
[ConvStack](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md).

> **Example**
> ```pycon
>>> import torch
>>> from luxonis_train.nodes.blocks import EncoderBlock
>>> x = torch.zeros(1, 3, 16, 16)
>>> EncoderBlock(3, 8, n_repeats=2)(x).shape
torch.Size([1, 8, 8, 8])
>>> EncoderBlock(3, 8, n_repeats=2, max_pool=False)(x).shape
torch.Size([1, 8, 16, 16])
```

#### Methods

##### init

```python
def __init__(in_channels: int, out_channels: int, n_repeats: int, *, max_pool: bool = True):
```

Build the pooling and the convolution stack.

Parameters

 * `in_channels` (`int`): The number of input channels.
 * `out_channels` (`int`): The number of output channels.
 * `n_repeats` (`int`): The number of `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) layers in the stack.
 * `max_pool` (`bool`): Whether a `2x2` max pool with stride `2` halves the size before the stack. Otherwise the step keeps the size.

### SimpleDecoder

Decoder of [SimpleDecoderBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md) steps, without skip connections.

The decoder takes the output of a [SimpleEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md) whose `base_hidden_channels` equals `base_width` and whose multipliers are the same. It restores the input size of the encoder when that size is divisible by 2n, where n is the number of multipliers. The example of [SimpleEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md) shows the pair.

#### Methods

##### init

```python
def __init__(base_width: int, out_channels: int, encoder_width_multipliers: list[float], n_convolutions: int = 2):
```

Build the decoder steps and the output convolution.

Parameters

 * `base_width` (`int`): The `base_hidden_channels` of the encoder.
 * `out_channels` (`int`): The number of output channels.
 * `encoder_width_multipliers` (`list[float]`): The `width_multipliers` of the encoder.
 * `n_convolutions` (`int`): The number of `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) layers in the stack of each step.
 * `block`: The class of the decoder steps.

##### forward

```python
def forward(x: Tensor) -> Tensor:
```

Run the decoder steps and the output convolution.

Parameters

 * `x` (`Tensor`): The output of [SimpleEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md), of shape `[B, C, h, w]`, with C = ⌊m⋅wb⌋. Here m is the last multiplier and wb is `base_width`.

Returns

 * `Tensor`: The output of shape `[B, out_channels, h * 2^n, w * 2^n]`, where `n` is the number of multipliers.

### SimpleDecoderBlock

Decoder step without a skip connection.

The step upsamples the input by `2` and runs the [ConvStack](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) on the result.

#### Methods

##### init

```python
def __init__(in_channels: int, out_channels: int, kernel_size: int, use_norm: bool, align_corners: bool, upsample_mode:
Literal['simple_upsample', 'conv_upsample', 'conv_transpose'], n_repeats: int):
```

Build the upsampling block and the convolution stack.

The stack maps `out_channels` channels to `out_channels`.

Parameters

 * `in_channels` (`int`): The number of input channels.
 * `out_channels` (`int`): The number of output channels.
 * `kernel_size` (`int`): The kernel size of [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md). Only `"conv_transpose"` uses it.
 * `use_norm` (`bool`): Whether the `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) of [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md) has a batch norm.
 * `align_corners` (`bool`): The `align_corners` option of the interpolation in [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
 * `upsample_mode` (`Literal['simple_upsample', 'conv_upsample', 'conv_transpose']`): The upsampling method of [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
 * `n_repeats` (`int`): The number of `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) layers in the stack.

##### forward

```python
def forward(x: Tensor) -> Tensor:
```

Upsample the input and run the convolution stack.

Parameters

 * `x` (`Tensor`): The input of shape `[B, in_channels, H, W]`.

Returns

 * `Tensor`: The output of shape `[B, out_channels, 2H, 2W]` for the interpolation modes. [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md) gives the size for `"conv_transpose"`.

### SimpleEncoder

Encoder of [EncoderBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md) steps that returns the last feature map.

The encoder has one step for each entry of `width_multipliers`, and one more step with the last entry again. A step with the entry m has ⌊m⋅c⌋ output channels, where c is `base_hidden_channels`. The first step keeps the size. Each later step halves it and rounds down. With n entries, the output is thus 2n times smaller than the input. [RecSubNet](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/backbones/recsubnet/recsubnet.md) pairs the encoder with [SimpleDecoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).

> **Example**
> ```pycon
>>> import torch
>>> from luxonis_train.nodes.blocks import (
...     SimpleDecoder,
...     SimpleEncoder,
... )
>>> encoder = SimpleEncoder(3, 8, [1, 2])
>>> len(encoder)
3
>>> features = encoder(torch.zeros(1, 3, 16, 16))
>>> features.shape
torch.Size([1, 16, 4, 4])
>>> SimpleDecoder(8, 3, [1, 2])(features).shape
torch.Size([1, 3, 16, 16])
```

#### Methods

##### init

```python
def __init__(in_channels: int, base_hidden_channels: int, width_multipliers: list[float], n_convolutions: int = 2):
```

Build the encoder steps.

Parameters

 * `in_channels` (`int`): The number of input channels.
 * `base_hidden_channels` (`int`): The base width. Each step multiplies it by its entry of `width_multipliers`.
 * `width_multipliers` (`list[float]`): The width factor of each step. An empty list makes the constructor raise `IndexError`.
 * `n_convolutions` (`int`): The number of `3x3`
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   layers in each step.

### UNetDecoder

Decoder of
[UNetDecoderBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
steps with skip connections.

The decoder takes the feature maps of a
[UNetEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
whose `base_hidden_channels` equals `base_width` and whose multipliers are the same.
[forward](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
removes the last map from the list, as the example shows.

> **Example**
> ```pycon
>>> import torch
>>> from luxonis_train.nodes.blocks import UNetDecoder, UNetEncoder
>>> encoder = UNetEncoder(3, 8, [1, 2])
>>> decoder = UNetDecoder(8, 2, [1, 2])
>>> features = encoder(torch.zeros(1, 3, 16, 16))
>>> decoder(features).shape
torch.Size([1, 2, 16, 16])
>>> len(features)
2
```

#### Methods

##### init

```python
def __init__(base_width: int, out_channels: int, encoder_width_multipliers: list[float], n_convolutions: int = 2):
```

Build the decoder steps and the output convolution.

Parameters

 * `base_width` (`int`): The `base_hidden_channels` of the encoder.
 * `out_channels` (`int`): The number of output channels.
 * `encoder_width_multipliers` (`list[float]`): The `width_multipliers` of the encoder.
 * `n_convolutions` (`int`): The number of `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) layers in the stack of each step.
 * `block`: The class of the decoder steps.

##### forward

```python
def forward(inputs: list[Tensor]) -> Tensor:
```

Decode the feature maps of [UNetEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).

The method pops the smallest map from `inputs` and starts from it. Each step then takes the next remaining map as its skip connection, from the smallest to the largest.

Warning: The method removes the last element of `inputs`.

Parameters

 * `inputs` (`list[Tensor]`): The feature maps of [UNetEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md), from the largest to the smallest. The list must hold one map more than the decoder has steps. Otherwise `zip` raises `ValueError`.

Returns

 * `Tensor`: The output of shape `[B, out_channels, H, W]`, where `H` and `W` are the size of the first map.

### UNetDecoderBlock

Decoder step with a skip connection from the encoder.

The step upsamples the input by `2` and concatenates the skip feature map along the channel axis. Then it runs the [ConvStack](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) on the result. The skip feature map must have `in_channels` channels and the size of the upsampled input.

> **Example**
> ```pycon
>>> import torch
>>> from luxonis_train.nodes.blocks import UNetDecoderBlock
>>> block = UNetDecoderBlock(
...     in_channels=8,
...     out_channels=4,
...     kernel_size=3,
...     use_norm=True,
...     align_corners=True,
...     upsample_mode="simple_upsample",
...     n_repeats=1,
... )
>>> x, skip_x = torch.zeros(1, 8, 4, 4), torch.zeros(1, 8, 8, 8)
>>> block(x, skip_x).shape
torch.Size([1, 4, 8, 8])
```

#### Methods

##### init

```python
def __init__(in_channels: int, out_channels: int, kernel_size: int, use_norm: bool, align_corners: bool, upsample_mode: Literal['simple_upsample', 'conv_upsample', 'conv_transpose'], n_repeats: int):
```

Build the upsampling block and the convolution stack.

The stack maps `in_channels + out_channels` channels to `out_channels`.

Parameters

 * `in_channels` (`int`): The number of channels of the input and of the skip feature map.
 * `out_channels` (`int`): The number of output channels.
 * `kernel_size` (`int`): The kernel size of
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
   Only `"conv_transpose"` uses it.
 * `use_norm` (`bool`): Whether the `3x3`
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   of
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
   has a batch norm.
 * `align_corners` (`bool`): The `align_corners` option of the interpolation in
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
 * `upsample_mode` (`Literal['simple_upsample', 'conv_upsample', 'conv_transpose']`): The upsampling method of
   [UpBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md).
 * `n_repeats` (`int`): The number of `3x3`
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   layers in the stack.

##### forward

```python
def forward(x: Tensor, skip_x: Tensor) -> Tensor:
```

Upsample `x`, concatenate `skip_x`, and run the stack.

Parameters

 * `x` (`Tensor`): The input of shape `[B, in_channels, H, W]`.
 * `skip_x` (`Tensor`): The encoder feature map of shape `[B, in_channels, 2H, 2W]`.

Returns

 * `Tensor`: The output of shape `[B, out_channels, 2H, 2W]`.

### UNetEncoder

[SimpleEncoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
that returns the feature map of every step.

[UNetDecoder](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/unet.md)
takes the list and uses the maps as skip connections.
[DiscSubNetHead](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/heads/discsubnet_head/discsubnet_head.md)
pairs the two.

> **Example**
> ```pycon
>>> import torch
>>> from luxonis_train.nodes.blocks import UNetEncoder
>>> encoder = UNetEncoder(3, 8, [1, 2])
>>> features = encoder(torch.zeros(1, 3, 16, 16))
>>> [tuple(feature.shape) for feature in features]
[(1, 8, 16, 16), (1, 16, 8, 8), (1, 16, 4, 4)]
```

#### Methods

##### forward

```python
def forward(x: Tensor) -> list[Tensor]:
```

Run the steps in order and collect every output.

Parameters

 * `x` (`Tensor`): The input of shape `[B, in_channels, H, W]`.

Returns

 * `list[Tensor]`: One feature map for each step, from the largest to the smallest. The map of step `i` has 2i times smaller height and width than `x`.

### UpBlock

Upsampling by `stride`, followed by a `3x3` [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md).

`upsample_mode` selects the upsampling:

 * `"conv_transpose"`: a `torch.nn.ConvTranspose2d` from `in_channels` to `out_channels`, with `kernel_size`, `stride`, and no padding. The output size is (H − 1)⋅s + k, with the stride s and the kernel size k.
 * `"simple_upsample"`: a `torch.nn.Upsample` by `stride`. The [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) then maps `in_channels` to `out_channels`.
 * `"conv_upsample"`: the same `torch.nn.Upsample`, then a `1x1` `torch.nn.Conv2d` from `in_channels` to `out_channels`.

The [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md) keeps the size and gives `out_channels` channels.

> **Example**
> ```pycon
>>> import torch
>>> from luxonis_train.nodes.blocks import UpBlock
>>> x = torch.zeros(1, 8, 4, 4)
>>> options = {
...     "kernel_size": 2,
...     "use_norm": True,
...     "align_corners": False,
... }
>>> UpBlock(8, 4, "conv_transpose", **options)(x).shape
torch.Size([1, 4, 8, 8])
>>> UpBlock(8, 4, "conv_upsample", **options)(x).shape
torch.Size([1, 4, 8, 8])
```

#### Methods

##### init

```python
def __init__(in_channels: int, out_channels: int, upsample_mode: Literal['simple_upsample', 'conv_upsample', 'conv_transpose'], kernel_size: int, use_norm: bool, align_corners: bool, stride: int = 2, activation: nn.Module | bool | None = True, interpolation_mode: Literal['nearest', 'linear', 'bilinear', 'bicubic', 'trilinear'] = 'bilinear'):
```

Build the upsampling layers and the convolution block.

Parameters

 * `in_channels` (`int`): The number of input channels.
 * `out_channels` (`int`): The number of output channels.
 * `upsample_mode` (`Literal['simple_upsample', 'conv_upsample', 'conv_transpose']`): The upsampling method. The class description
   lists the layers of each method.
 * `kernel_size` (`int`): The kernel size of the transposed convolution. The other methods ignore it.
 * `use_norm` (`bool`): Whether the
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md)
   has a batch norm.
 * `align_corners` (`bool`): The `align_corners` option of `torch.nn.Upsample`. `"conv_transpose"` ignores it.
 * `stride` (`int`): The upsampling factor.
 * `activation` (`nn.Module | bool | None`): The activation of the
   [ConvBlock](https://docs.luxonis.com/software-v3/ai-inference/model-source/training/luxonis-train/luxonis-train-api-reference/nodes/blocks/blocks.md).
   `True` selects `torch.nn.ReLU`. `False` or `None` selects `torch.nn.Identity`.
 * `interpolation_mode` (`Literal['nearest', 'linear', 'bilinear', 'bicubic', 'trilinear']`): The mode of `torch.nn.Upsample`.
   `"conv_transpose"` ignores it. Warning: only `"bilinear"` and `"bicubic"` work on a 4D input. `"nearest"` rejects any
   `align_corners` value. `"linear"` needs a 3D input, and `"trilinear"` needs a 5D input.