| | """Contains the UNet decoder. |
| | |
| | For licensing see accompanying LICENSE file. |
| | Copyright (C) 2025 Apple Inc. All Rights Reserved. |
| | """ |
| |
|
| | from __future__ import annotations |
| |
|
| | from typing import List |
| |
|
| | import torch |
| | from torch import nn |
| |
|
| | from sharp.models.blocks import ( |
| | NormLayerName, |
| | norm_layer_2d, |
| | residual_block_2d, |
| | ) |
| |
|
| | from .base_decoder import BaseDecoder |
| |
|
| |
|
| | class UNetDecoder(BaseDecoder): |
| | """Decoder of UNet model.""" |
| |
|
| | def __init__( |
| | self, |
| | dim_out: int, |
| | width: List[int] | int, |
| | steps: int = 5, |
| | norm_type: NormLayerName = "group_norm", |
| | norm_num_groups=8, |
| | blocks_per_layer=2, |
| | ) -> None: |
| | """Initialize UNet Decoder. |
| | |
| | Args: |
| | dim_out: The number of output channels. |
| | width: Width of last input feature map from encoder |
| | or the width list of all input feature maps from encoder. |
| | steps: The number of upsampling steps. |
| | norm_type: Which kind of normalization layer to use. |
| | norm_num_groups: How many groups to use for group norm (if relevant). |
| | blocks_per_layer: How many blocks per layer to use. |
| | """ |
| | super().__init__() |
| |
|
| | if blocks_per_layer < 1: |
| | raise ValueError("blocks_per_layer must be greater or equal to one.") |
| |
|
| | self.dim_out = dim_out |
| |
|
| | self.convs_up = nn.ModuleList() |
| |
|
| | self.output_dims: list[int] |
| | |
| | if isinstance(width, int): |
| | self.input_dims = [width >> i for i in range(0, steps + 1)] |
| | else: |
| | self.input_dims = width[::-1][: steps + 1] |
| |
|
| | for i_step in range(steps): |
| | input_width = self.input_dims[i_step] |
| | current_width = self.input_dims[i_step + 1] |
| | convs_up_i = nn.Sequential( |
| | nn.Upsample(scale_factor=2), |
| | residual_block_2d( |
| | input_width * (1 if i_step == 0 else 2), |
| | current_width, |
| | norm_type=norm_type, |
| | norm_num_groups=norm_num_groups, |
| | ), |
| | *[ |
| | residual_block_2d( |
| | current_width, |
| | current_width, |
| | norm_type=norm_type, |
| | norm_num_groups=norm_num_groups, |
| | ) |
| | for _ in range(blocks_per_layer - 1) |
| | ], |
| | ) |
| | self.convs_up.append(convs_up_i) |
| | input_width = 2 * current_width |
| | current_width //= 2 |
| |
|
| | last_width = self.input_dims[-1] |
| | self.conv_out = nn.Sequential( |
| | norm_layer_2d(last_width * 2, norm_type, num_groups=norm_num_groups), |
| | nn.ReLU(), |
| | nn.Conv2d(last_width * 2, dim_out, 1), |
| | norm_layer_2d(dim_out, norm_type, num_groups=norm_num_groups), |
| | nn.ReLU(), |
| | ) |
| |
|
| | def forward(self, features: list[torch.Tensor]) -> torch.Tensor: |
| | """Apply UNet to image. |
| | |
| | Args: |
| | features: The input multi-level feature map from encoder. |
| | |
| | Returns: |
| | The output feature map. |
| | """ |
| | i_feature_layer = len(features) - 1 |
| | out = self.convs_up[0](features[i_feature_layer]) |
| | i_feature_layer -= 1 |
| | for conv_up in self.convs_up[1:]: |
| | out = conv_up(torch.cat([out, features[i_feature_layer]], dim=1)) |
| | i_feature_layer -= 1 |
| | out = self.conv_out(torch.cat([out, features[i_feature_layer]], dim=1)) |
| |
|
| | return out |
| |
|
