Method __init__

src/sharp/models/blocks.py:134–170

Initialize feature fusion block. Args: dim_in: Dimensions of input. dim_out: Dimensions of output. batch_norm: Whether to use batch normalization in resnet blocks. upsampling_mode: What mode to use for upsampling. None if no upsampling

(
        self,
        dim_in: int,
        dim_out: int | None = None,
        upsampling_mode: UpsamplingMode | None = None,
        batch_norm: bool = False,
    )

    deconv: nn.Module

    def __init__(
        self,
        dim_in: int,
        dim_out: int | None = None,
        upsampling_mode: UpsamplingMode | None = None,
        batch_norm: bool = False,
    ):
        """Initialize feature fusion block.

        Args:
            dim_in: Dimensions of input.
            dim_out: Dimensions of output.
            batch_norm: Whether to use batch normalization in resnet blocks.
            upsampling_mode: What mode to use for upsampling. None if no upsampling
                is required.
        """
        super().__init__()
        if dim_out is None:
            dim_out = dim_in
        self.resnet1 = self._residual_block(dim_in, batch_norm)
        self.resnet2 = self._residual_block(dim_in, batch_norm)

        if upsampling_mode is not None:
            self.deconv = upsampling_layer(upsampling_mode, scale_factor=2, dim_in=dim_in)
        else:
            self.deconv = nn.Sequential()

        self.out_conv = nn.Conv2d(
            dim_in,
            dim_out,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=True,
        )

        self.skip_add = nn.quantized.FloatFunctional()

    def forward(self, x0: torch.Tensor, x1: torch.Tensor | None = None) -> torch.Tensor:
        """Process and fuse input features."""