Function _create_project_upsample_block

src/sharp/models/gaussian_decoder.py:50–83

(
    dim_in: int,
    dim_out: int,
    upsample_layers: int,
    dim_intermediate: int | None = None,
)

def _create_project_upsample_block(
    dim_in: int,
    dim_out: int,
    upsample_layers: int,
    dim_intermediate: int | None = None,
) -> nn.Module:
    if dim_intermediate is None:
        dim_intermediate = dim_out
    # Projection.
    blocks = [
        nn.Conv2d(
            in_channels=dim_in,
            out_channels=dim_intermediate,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=False,
        )
    ]

    # Upsampling.
    blocks += [
        nn.ConvTranspose2d(
            in_channels=dim_intermediate if i == 0 else dim_out,
            out_channels=dim_out,
            kernel_size=2,
            stride=2,
            padding=0,
            bias=False,
        )
        for i in range(upsample_layers)
    ]

    return nn.Sequential(*blocks)


class ImageFeatures(NamedTuple):