diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d27d16547..6afbdf7f4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,7 +17,7 @@ repos: rev: v0.16.1 hooks: - id: ruff - exclude: "^(docs/.*|src/careamics/lvae_training/.*|src/careamics/models/lvae/(?!noise_models\\.py).*|scripts/.*|demos/.*)" + exclude: "^(docs/.*|src/careamics/lvae_training/.*|src/careamics/models/lvae/noise_models.py|scripts/.*|demos/.*)" - repo: https://github.com/psf/black-pre-commit-mirror rev: 26.5.1 @@ -30,7 +30,7 @@ repos: hooks: - id: mypy files: "^src/" - exclude: "^(docs/.*|src/careamics/lvae_training/.*|src/careamics/models/lvae/(?!noise_models\\.py).*|src/careamics/config/likelihood_model.py|src/careamics/losses/lvae/.*)" + exclude: "^(docs/.*|src/careamics/lvae_training/.*|src/careamics/models/lvae/noise_models.py|src/careamics/config/likelihood_model.py|src/careamics/losses/lvae/.*)" args: ["--config-file", "mypy.ini"] additional_dependencies: - numpy @@ -44,7 +44,7 @@ repos: rev: v1.11.0rc0 hooks: - id: numpydoc-validation - exclude: "^(tests/.*|docs/.*|src/careamics/lvae_training/.*|src/careamics/models/lvae/(?!noise_models\\.py).*|src/careamics/losses/lvae/.*|scripts/.*)" + exclude: "^(tests/.*|docs/.*|src/careamics/lvae_training/.*|src/careamics/models/lvae/noise_models.py|src/careamics/losses/lvae/.*|scripts/.*)" # data_module: overloads and config params (val_percentage etc.) don't match signature; see https://github.com/numpy/numpydoc/issues/559 # # jupyter linting and formatting diff --git a/src/careamics/config/architectures/lvae_config.py b/src/careamics/config/architectures/lvae_config.py index f0f3ee865..b2bd34274 100644 --- a/src/careamics/config/architectures/lvae_config.py +++ b/src/careamics/config/architectures/lvae_config.py @@ -142,20 +142,6 @@ def validate_z_dims(cls, z_dims: tuple) -> tuple: return z_dims - def set_3D(self, is_3D: bool) -> None: - """ - Set 3D model by setting the `conv_dims` parameters. - - Parameters - ---------- - is_3D : bool - Whether the algorithm is 3D or not. - """ - if is_3D: - self.conv_dims = 3 - else: - self.conv_dims = 2 - def is_3D(self) -> bool: """ Return whether the model is 3D or not. diff --git a/src/careamics/models/lvae/__init__.py b/src/careamics/models/lvae/__init__.py index 4b0d7af22..cf03f0d7b 100644 --- a/src/careamics/models/lvae/__init__.py +++ b/src/careamics/models/lvae/__init__.py @@ -1,3 +1,5 @@ +"""LVAE model package.""" + __all__ = ["LadderVAE"] from .lvae import LadderVAE diff --git a/src/careamics/models/lvae/layers.py b/src/careamics/models/lvae/layers.py index 7665d2188..19a70f925 100644 --- a/src/careamics/models/lvae/layers.py +++ b/src/careamics/models/lvae/layers.py @@ -1,8 +1,8 @@ """Script containing the common basic blocks (nn.Module) reused by the LadderVAE.""" -from collections.abc import Iterable +from collections.abc import Callable, Iterable, Sequence from copy import deepcopy -from typing import Callable, Literal, Optional, Union +from typing import Literal, Union import numpy as np import torch @@ -18,27 +18,38 @@ NormType = Union[nn.BatchNorm2d, nn.BatchNorm3d] DropoutType = Union[nn.Dropout2d, nn.Dropout3d] +_DEFAULT_NONLIN = nn.LeakyReLU() + class ResidualBlock(nn.Module): """ Residual block with 2 convolutional layers. + The block follows the fixed "bacdbacd" structure used by the LadderVAE, namely two + repetitions of [batchnorm, activation, conv, dropout], optionally followed by a + gating layer. The output is given by: ``out = [gate](f(x)) + x``. + Some architectural notes: - The number of input, intermediate, and output channels is the same, - Padding is always 'same', - The 2 convolutional layers have the same groups, - No stride allowed, - - Kernel sizes must be odd. - - The output isgiven by: `out = gate(f(x)) + x`. - The presence of the gating mechanism is optional, and f(x) has different - structures depending on the `block_type` argument. - Specifically, `block_type` is a string specifying the block's structure, with: - a = activation - b = batch norm - c = conv layer - d = dropout. - For example, "bacdbacd" defines a block with 2x[batchnorm, activation, conv, dropout]. + - Kernel size is fixed to 3. + + Parameters + ---------- + channels : int + The number of input and output channels (they are the same). + nonlin : Callable + The non-linearity function used in the block (e.g., `nn.ReLU`). + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + groups : int, optional + The number of groups to consider in the convolutions. Default is 1. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + gated : bool, optional + Whether to append a gating layer at the end of the block. Default is `False`. """ default_kernel_size = (3, 3) @@ -47,114 +58,56 @@ def __init__( self, channels: int, nonlin: Callable, - conv_strides: tuple[int] = (2, 2), - kernel: Union[int, Iterable[int], None] = None, + conv_strides: Sequence[int] = (2, 2), groups: int = 1, - batchnorm: bool = True, - block_type: str = None, - dropout: float = None, - gated: bool = None, - conv2d_bias: bool = True, + dropout: float | None = None, + gated: bool = False, ): """ Constructor. Parameters ---------- - channels: int + channels : int The number of input and output channels (they are the same). - nonlin: Callable + nonlin : Callable The non-linearity function used in the block (e.g., `nn.ReLU`). - kernel: Union[int, Iterable[int]], optional - The kernel size used in the convolutions of the block. - It can be either a single integer or a pair of integers defining the squared kernel. - Default is `None`. - groups: int, optional + conv_strides : tuple of int, optional + The convolution strides, used to infer the convolution dimensionality. + Default is `(2, 2)`. + groups : int, optional The number of groups to consider in the convolutions. Default is 1. - batchnorm: bool, optional - Whether to use batchnorm layers. Default is `True`. - block_type: str, optional - A string specifying the block structure, check class docstring for more info. - Default is `None`. - dropout: float, optional - The dropout probability in dropout layers. If `None` dropout is not used. - Default is `None`. - gated: bool, optional - Whether to use gated layer. Default is `None`. - conv2d_bias: bool, optional - Whether to use bias term in convolutions. Default is `True`. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + gated : bool, optional + Whether to append a gating layer at the end of the block. Default is + `False`. """ super().__init__() - # Set kernel size & padding - if kernel is None: - kernel = self.default_kernel_size - elif isinstance(kernel, int): - kernel = (kernel, kernel) - elif len(kernel) != 2: - raise ValueError("kernel has to be None, int, or an iterable of length 2") - assert all(k % 2 == 1 for k in kernel), "kernel sizes have to be odd" - kernel = list(kernel) - - # Define modules + kernel = list(self.default_kernel_size) + conv_layer: ConvType = getattr(nn, f"Conv{len(conv_strides)}d") norm_layer: NormType = getattr(nn, f"BatchNorm{len(conv_strides)}d") dropout_layer: DropoutType = getattr(nn, f"Dropout{len(conv_strides)}d") - # TODO: same comment as in lvae.py, would be more readable to have `conv_dims` + # "bacdbacd" block: 2x [batchnorm, activation, conv, dropout] modules = [] - if block_type == "cabdcabd": - for i in range(2): - conv = conv_layer( - channels, - channels, - kernel[i], - padding="same", - groups=groups, - bias=conv2d_bias, - ) - modules.append(conv) - modules.append(nonlin) - if batchnorm: - modules.append(norm_layer(channels)) - if dropout is not None: - modules.append(dropout_layer(dropout)) - elif block_type == "bacdbac": - for i in range(2): - if batchnorm: - modules.append(norm_layer(channels)) - modules.append(nonlin) - conv = conv_layer( - channels, - channels, - kernel[i], - padding="same", - groups=groups, - bias=conv2d_bias, - ) - modules.append(conv) - if dropout is not None and i == 0: - modules.append(dropout_layer(dropout)) - elif block_type == "bacdbacd": - for i in range(2): - if batchnorm: - modules.append(norm_layer(channels)) - modules.append(nonlin) - conv = conv_layer( + for i in range(2): + modules.append(norm_layer(channels)) + modules.append(nonlin) + modules.append( + conv_layer( channels, channels, kernel[i], padding="same", groups=groups, - bias=conv2d_bias, ) - modules.append(conv) - modules.append(dropout_layer(dropout)) - - else: - raise ValueError(f"unrecognized block type '{block_type}'") + ) + modules.append(dropout_layer(dropout)) - self.gated = gated + # Optional gating mechanism if gated: modules.append( GateLayer( @@ -173,12 +126,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Parameters ---------- x : torch.Tensor - input tensor # TODO add shape + Input tensor of shape (B, C, [Z], Y, X). Returns ------- torch.Tensor - output tensor # TODO add shape + Output tensor of the same shape as the input. """ out = self.block(x) assert ( @@ -187,28 +140,47 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return out + x -class ResidualGatedBlock(ResidualBlock): - """Layer class that implements a residual block with a gating mechanism.""" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs, gated=True) - - class GateLayer(nn.Module): """ Layer class that implements a gating mechanism. Double the number of channels through a convolutional layer, then use half the channels as gate for the other half. + + Parameters + ---------- + channels : int + The number of input (and output) channels. + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + Default is `(2, 2)`. + kernel_size : int, optional + The size of the convolution kernel. Default is 3. + nonlin : Callable, optional + The non-linearity applied to the non-gate half. Default is `nn.LeakyReLU`. """ def __init__( self, channels: int, - conv_strides: tuple[int] = (2, 2), + conv_strides: Sequence[int] = (2, 2), kernel_size: int = 3, - nonlin: Callable = nn.LeakyReLU(), + nonlin: Callable = _DEFAULT_NONLIN, ): + """Constructor. + + Parameters + ---------- + channels : int + The number of input (and output) channels. + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + Default is `(2, 2)`. + kernel_size : int, optional + The size of the convolution kernel. Default is 3. + nonlin : Callable, optional + The non-linearity applied to the non-gate half. Default is `nn.LeakyReLU`. + """ super().__init__() assert kernel_size % 2 == 1 pad = kernel_size // 2 @@ -222,200 +194,269 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Parameters ---------- x : torch.Tensor - input # TODO add shape + Input tensor of shape (B, C, [Z], Y, X). Returns ------- torch.Tensor - output # TODO add shape + The gated output tensor of shape (B, C, [Z], Y, X). """ x = self.conv(x) x, gate = torch.chunk(x, 2, dim=1) - x = self.nonlin(x) # TODO remove this? + x = self.nonlin(x) gate = torch.sigmoid(gate) return x * gate -class ResBlockWithResampling(nn.Module): +def _make_pre_conv( + direction: Literal["top-down", "bottom-up"], + c_in: int, + c_out: int, + conv_strides: Sequence[int], + resample: bool, + groups: int, +) -> Union[nn.Module, None]: """ - Residual block with resampling. - - Residual block that takes care of resampling (i.e. downsampling or upsampling) steps (by a factor 2). - It is structured as follows: - 1. `pre_conv`: a downsampling or upsampling strided convolutional layer in case of resampling, or - a 1x1 convolutional layer that maps the number of channels of the input to `inner_channels`. - 2. `ResidualBlock` - 3. `post_conv`: a 1x1 convolutional layer that maps the number of channels to `c_out`. - - Some implementation notes: - - Resampling is performed through a strided convolution layer at the beginning of the block. - - The strided convolution block has fixed kernel size of 3x3 and 1 layer of padding with zeros. - - The number of channels is adjusted at the beginning and end of the block through 1x1 convolutional layers. - - The number of internal channels is by default the same as the number of output channels, but - min_inner_channels can override the behaviour. + Build the input convolution of a deterministic resampling residual block. + + The convolution performs (optional) resampling by a factor 2 and/or maps the input + to `c_out` channels: + - if `resample` is `True`: a strided (transposed) convolution that downsamples + ("bottom-up") or upsamples ("top-down") by a factor 2, + - elif `c_in != c_out`: a 1x1 convolution mapping the channels, + - else: `None` (no input convolution needed). + + Parameters + ---------- + direction : Literal["top-down", "bottom-up"] + The resampling direction. "bottom-up" downsamples, "top-down" upsamples. + c_in : int + The number of input channels. + c_out : int + The number of output channels. + conv_strides : tuple of int + The convolution strides, used to infer the convolution dimensionality. + resample : bool + Whether to resample (by a factor 2) in this convolution. + groups : int + The number of groups to consider in the convolution. + + Returns + ------- + torch.nn.Module or None + The input convolution, or `None` if no channel change nor resampling is + required. + """ + conv_layer: ConvType = getattr(nn, f"Conv{len(conv_strides)}d") + + if resample: + if direction == "bottom-up": # downsample + return conv_layer( + in_channels=c_in, + out_channels=c_out, + kernel_size=3, + padding=1, + stride=conv_strides, + groups=groups, + ) + # top-down: upsample + transp_conv_layer: ConvType = getattr(nn, f"ConvTranspose{len(conv_strides)}d") + return transp_conv_layer( + in_channels=c_in, + out_channels=c_out, + kernel_size=3, + padding=1, + stride=conv_strides, + groups=groups, + output_padding=1 if len(conv_strides) == 2 else (0, 1, 1), + ) + if c_in != c_out: + return conv_layer(c_in, c_out, 1, groups=groups) + return None + + +class BottomUpDeterministicResBlock(nn.Module): + """ + Resnet block for bottom-up (downsampling) deterministic layers. + + It is structured as an (optional) downsampling `pre_conv` strided convolution + followed by a `ResidualBlock`. + + Parameters + ---------- + c_in : int + The number of input channels. + c_out : int + The number of output channels. + conv_strides : Sequence[int] + The convolution strides, used to infer the convolution dimensionality. + nonlin : Callable, optional + The non-linearity function used in the block. Default is `nn.LeakyReLU`. + downsample : bool, optional + Whether to downsample by a factor 2 in `pre_conv`. Default is `False`. + groups : int, optional + The number of groups to consider in the convolutions. Default is 1. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + gated : bool, optional + Whether to use a gated residual block. Default is `False`. """ def __init__( self, - mode: Literal["top-down", "bottom-up"], c_in: int, c_out: int, - conv_strides: tuple[int], - min_inner_channels: Union[int, None] = None, - nonlin: Callable = nn.LeakyReLU(), - resample: bool = False, - res_block_kernel: Optional[Union[int, Iterable[int]]] = None, + conv_strides: Sequence[int], + nonlin: Callable = _DEFAULT_NONLIN, + downsample: bool = False, groups: int = 1, - batchnorm: bool = True, - res_block_type: Union[str, None] = None, dropout: Union[float, None] = None, - gated: Union[bool, None] = None, - conv2d_bias: bool = True, - # lowres_input: bool = False, + gated: bool = False, ): """ Constructor. Parameters ---------- - mode: Literal["top-down", "bottom-up"] - The type of resampling performed in the initial strided convolution of the block. - If "bottom-up" downsampling of a factor 2 is done. - If "top-down" upsampling of a factor 2 is done. - c_in: int + c_in : int The number of input channels. - c_out: int + c_out : int The number of output channels. - min_inner_channels: int, optional - The number of channels used in the inner layer of this module. - Default is `None`, meaning that the number of inner channels is set to `c_out`. - nonlin: Callable, optional + conv_strides : tuple of int + The convolution strides, used to infer the convolution dimensionality. + nonlin : Callable, optional The non-linearity function used in the block. Default is `nn.LeakyReLU`. - resample: bool, optional - Whether to perform resampling in the first convolutional layer. - If `False`, the first convolutional layer just maps the input to a tensor with - `inner_channels` channels through 1x1 convolution. Default is `False`. - res_block_kernel: Union[int, Iterable[int]], optional - The kernel size used in the convolutions of the residual block. - It can be either a single integer or a pair of integers defining the squared kernel. - Default is `None`. - groups: int, optional + downsample : bool, optional + Whether to downsample by a factor 2 in the input convolution. Default is + `False`. + groups : int, optional The number of groups to consider in the convolutions. Default is 1. - batchnorm: bool, optional - Whether to use batchnorm layers. Default is `True`. - res_block_type: str, optional - A string specifying the structure of residual block. - Check `ResidualBlock` doscstring for more information. - Default is `None`. - dropout: float, optional - The dropout probability in dropout layers. If `None` dropout is not used. - Default is `None`. - gated: bool, optional - Whether to use gated layer. Default is `None`. - conv2d_bias: bool, optional - Whether to use bias term in convolutions. Default is `True`. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + gated : bool, optional + Whether to use a gated residual block. Default is `False`. """ super().__init__() - assert mode in ["top-down", "bottom-up"] - - conv_layer: ConvType = getattr(nn, f"Conv{len(conv_strides)}d") - transp_conv_layer: ConvType = getattr(nn, f"ConvTranspose{len(conv_strides)}d") - - if min_inner_channels is None: - min_inner_channels = 0 - # inner_channels is the number of channels used in the inner layers - # of ResBlockWithResampling - inner_channels = max(c_out, min_inner_channels) - - # Define first conv layer to change num channels and/or up/downsample - if resample: - if mode == "bottom-up": # downsample - self.pre_conv = conv_layer( - in_channels=c_in, - out_channels=inner_channels, - kernel_size=3, - padding=1, - stride=conv_strides, - groups=groups, - bias=conv2d_bias, - ) - elif mode == "top-down": # upsample - self.pre_conv = transp_conv_layer( - in_channels=c_in, - kernel_size=3, - out_channels=inner_channels, - padding=1, # TODO maybe don't hardcode this? - stride=conv_strides, - groups=groups, - output_padding=1 if len(conv_strides) == 2 else (0, 1, 1), - bias=conv2d_bias, - ) - elif c_in != inner_channels: - self.pre_conv = conv_layer( - c_in, inner_channels, 1, groups=groups, bias=conv2d_bias - ) - else: - self.pre_conv = None - - # Residual block + self.pre_conv = _make_pre_conv( + "bottom-up", c_in, c_out, conv_strides, downsample, groups + ) self.res = ResidualBlock( - channels=inner_channels, + channels=c_out, conv_strides=conv_strides, nonlin=nonlin, - kernel=res_block_kernel, groups=groups, - batchnorm=batchnorm, dropout=dropout, gated=gated, - block_type=res_block_type, - conv2d_bias=conv2d_bias, ) - # Define last conv layer to get correct num output channels - if inner_channels != c_out: - self.post_conv = conv_layer( - inner_channels, c_out, 1, groups=groups, bias=conv2d_bias - ) - else: - self.post_conv = None - def forward(self, x: torch.Tensor) -> torch.Tensor: """Forward pass. Parameters ---------- x : torch.Tensor - input # TODO add shape + Input tensor of shape (B, C_in, [Z], Y, X). Returns ------- torch.Tensor - output # TODO add shape + Output tensor of shape (B, C_out, [Z'], Y', X'). """ if self.pre_conv is not None: x = self.pre_conv(x) + return self.res(x) - x = self.res(x) - - if self.post_conv is not None: - x = self.post_conv(x) - return x +class TopDownDeterministicResBlock(nn.Module): + """ + Resnet block for top-down (upsampling) deterministic layers. + + It is structured as an (optional) upsampling `pre_conv` transposed convolution + followed by a `ResidualBlock`. + + Parameters + ---------- + c_in : int + The number of input channels. + c_out : int + The number of output channels. + conv_strides : Sequence[int] + The convolution strides, used to infer the convolution dimensionality. + nonlin : Callable, optional + The non-linearity function used in the block. Default is `nn.LeakyReLU`. + upsample : bool, optional + Whether to upsample by a factor 2 in `pre_conv`. Default is `False`. + groups : int, optional + The number of groups to consider in the convolutions. Default is 1. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + gated : bool, optional + Whether to use a gated residual block. Default is `False`. + """ -class TopDownDeterministicResBlock(ResBlockWithResampling): - """Resnet block for top-down deterministic layers.""" + def __init__( + self, + c_in: int, + c_out: int, + conv_strides: Sequence[int], + nonlin: Callable = _DEFAULT_NONLIN, + upsample: bool = False, + groups: int = 1, + dropout: Union[float, None] = None, + gated: bool = False, + ): + """ + Constructor. - def __init__(self, *args, upsample: bool = False, **kwargs): - kwargs["resample"] = upsample - super().__init__("top-down", *args, **kwargs) + Parameters + ---------- + c_in : int + The number of input channels. + c_out : int + The number of output channels. + conv_strides : tuple of int + The convolution strides, used to infer the convolution dimensionality. + nonlin : Callable, optional + The non-linearity function used in the block. Default is `nn.LeakyReLU`. + upsample : bool, optional + Whether to upsample by a factor 2 in the input convolution. Default is + `False`. + groups : int, optional + The number of groups to consider in the convolutions. Default is 1. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + gated : bool, optional + Whether to use a gated residual block. Default is `False`. + """ + super().__init__() + self.pre_conv = _make_pre_conv( + "top-down", c_in, c_out, conv_strides, upsample, groups + ) + self.res = ResidualBlock( + channels=c_out, + conv_strides=conv_strides, + nonlin=nonlin, + groups=groups, + dropout=dropout, + gated=gated, + ) + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass. -class BottomUpDeterministicResBlock(ResBlockWithResampling): - """Resnet block for bottom-up deterministic layers.""" + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (B, C_in, [Z], Y, X). - def __init__(self, *args, downsample: bool = False, **kwargs): - kwargs["resample"] = downsample - super().__init__("bottom-up", *args, **kwargs) + Returns + ------- + torch.Tensor + Output tensor of shape (B, C_out, [Z'], Y', X'). + """ + if self.pre_conv is not None: + x = self.pre_conv(x) + return self.res(x) class BottomUpLayer(nn.Module): @@ -423,83 +464,114 @@ class BottomUpLayer(nn.Module): Bottom-up deterministic layer. It consists of one or a stack of `BottomUpDeterministicResBlock`'s. - The outputs are the so-called `bu_values` that are later used in the Decoder to update the + The outputs are the so-called `bu_values` that are later used in the Decoder to + update the generative distributions. NOTE: When Lateral Contextualization is Enabled (i.e., `enable_multiscale=True`), the low-res lateral input is first fed through a BottomUpDeterministicBlock (BUDB) - (without downsampling), and then merged to the latent tensor produced by the primary flow - of the `BottomUpLayer` through the `MergeLowRes` layer. It is meaningful to remark that + (without downsampling), and then merged to the latent tensor produced by the primary + flow + of the `BottomUpLayer` through the `MergeLowRes` layer. It is meaningful to remark + that the BUDB that takes care of encoding the low-res input can be either shared with the - primary flow (and in that case it is the "same_size" BUDB (or stack of BUDBs) -> see `self.net`), + primary flow (and in that case it is the "same_size" BUDB (or stack of BUDBs) -> see + `self.net`), or can be a deep-copy of the primary flow's BUDB. This behaviour is controlled by `lowres_separate_branch` parameter. + + Parameters + ---------- + n_res_blocks : int + Number of `BottomUpDeterministicResBlock` modules stacked in this layer. + n_filters : int + Number of channels present throughout the layers of this block. + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + downsampling_steps : int, optional + Number of downsampling steps done in this layer (typically 1). Default is 0. + nonlin : Callable, optional + The non-linearity function used in the block. Default is `nn.LeakyReLU`. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + enable_multiscale : bool, optional + Whether to enable multiscale (Lateral Contextualization). Default is `False`. + multiscale_lowres_size_factor : int, optional + Factor expressing the relative size of the primary-flow tensor with respect to + the lower-resolution lateral input tensor. Default is `None`. + lowres_separate_branch : bool, optional + Whether the low-res residual block(s) are shared (`False`) or not (`True`) with + the primary-flow "same-size" residual block(s). Default is `False`. + multiscale_retain_spatial_dims : bool, optional + Whether to pad the primary-flow latent to match the low-res input size. + Default is `False`. + decoder_retain_spatial_dims : bool, optional + Whether the corresponding top-down layer retains the spatial dims. Default + is `False`. + output_expected_shape : Iterable[int], optional + The expected output shape (only used if `enable_multiscale == True`). + Default is `None`. """ def __init__( self, n_res_blocks: int, n_filters: int, - conv_strides: tuple[int] = (2, 2), + conv_strides: Sequence[int] = (2, 2), downsampling_steps: int = 0, - nonlin: Optional[Callable] = None, - batchnorm: bool = True, - dropout: Optional[float] = None, - res_block_type: Optional[str] = None, - res_block_kernel: Optional[int] = None, - gated: Optional[bool] = None, + nonlin: Callable = _DEFAULT_NONLIN, + dropout: float | None = None, enable_multiscale: bool = False, - multiscale_lowres_size_factor: Optional[int] = None, + multiscale_lowres_size_factor: int | None = None, lowres_separate_branch: bool = False, multiscale_retain_spatial_dims: bool = False, decoder_retain_spatial_dims: bool = False, - output_expected_shape: Optional[Iterable[int]] = None, + output_expected_shape: Sequence[int] | None = None, ): """ Constructor. Parameters ---------- - n_res_blocks: int + n_res_blocks : int Number of `BottomUpDeterministicResBlock` modules stacked in this layer. - n_filters: int + n_filters : int Number of channels present through out the layers of this block. - downsampling_steps: int, optional - Number of downsampling steps that has to be done in this layer (typically 1). + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + Default is `(2, 2)`. + downsampling_steps : int, optional + Number of downsampling steps that has to be done in this layer (typically + 1). Default is 0. - nonlin: Callable, optional + nonlin : Callable, optional The non-linearity function used in the block. Default is `None`. - batchnorm: bool, optional - Whether to use batchnorm layers. Default is `True`. - dropout: float, optional + dropout : float, optional The dropout probability in dropout layers. If `None` dropout is not used. Default is `None`. - res_block_type: str, optional - A string specifying the structure of residual block. - Check `ResidualBlock` doscstring for more information. - Default is `None`. - res_block_kernel: Union[int, Iterable[int]], optional - The kernel size used in the convolutions of the residual block. - It can be either a single integer or a pair of integers defining the squared kernel. - Default is `None`. - gated: bool, optional - Whether to use gated layer. Default is `None`. - enable_multiscale: bool, optional - Whether to enable multiscale (Lateral Contextualization) or not. Default is `False`. - multiscale_lowres_size_factor: int, optional - A factor the expresses the relative size of the primary flow tensor with respect to the + enable_multiscale : bool, optional + Whether to enable multiscale (Lateral Contextualization) or not. Default is + `False`. + multiscale_lowres_size_factor : int, optional + A factor the expresses the relative size of the primary flow tensor with + respect to the lower-resolution lateral input tensor. Default in `None`. - lowres_separate_branch: bool, optional - Whether the residual block(s) encoding the low-res input should be shared (`False`) or - not (`True`) with the primary flow "same-size" residual block(s). Default is `False`. - multiscale_retain_spatial_dims: bool, optional - Whether to pad the latent tensor resulting from the bottom-up layer's primary flow + lowres_separate_branch : bool, optional + Whether the residual block(s) encoding the low-res input should be shared + (`False`) or + not (`True`) with the primary flow "same-size" residual block(s). Default is + `False`. + multiscale_retain_spatial_dims : bool, optional + Whether to pad the latent tensor resulting from the bottom-up layer's + primary flow to match the size of the low-res input. Default is `False`. - decoder_retain_spatial_dims: bool, optional - Whether in the corresponding top-down layer the shape of tensor is retained between + decoder_retain_spatial_dims : bool, optional + Whether in the corresponding top-down layer the shape of tensor is retained + between input and output. Default is `False`. - output_expected_shape: Iterable[int], optional - The expected shape of the layer output (only used if `enable_multiscale == True`). + output_expected_shape : Iterable[int], optional + The expected shape of the layer output (only used if `enable_multiscale == + True`). Default is `None`. """ super().__init__() @@ -526,11 +598,8 @@ def __init__( c_out=n_filters, nonlin=nonlin, downsample=do_resample, - batchnorm=batchnorm, dropout=dropout, - res_block_type=res_block_type, - res_block_kernel=res_block_kernel, - gated=gated, + gated=True, ) if do_resample: bu_blocks_downsized.append(block) @@ -541,30 +610,22 @@ def __init__( self.net = nn.Sequential(*bu_blocks_samesize) # Using the same net for the low resolution (and larger sized image) - self.lowres_net = self.lowres_merge = None + self.lowres_net: nn.Module | None = None + self.lowres_merge: nn.Module | None = None if self.enable_multiscale: self._init_multiscale( n_filters=n_filters, conv_strides=conv_strides, nonlin=nonlin, - batchnorm=batchnorm, dropout=dropout, - res_block_type=res_block_type, ) - # msg = f'[{self.__class__.__name__}] McEnabled:{int(enable_multiscale)} ' - # if enable_multiscale: - # msg += f'McParallelBeam:{int(multiscale_retain_spatial_dims)} McFactor{multiscale_lowres_size_factor}' - # print(msg) - def _init_multiscale( self, - nonlin: Callable = None, - n_filters: int = None, - conv_strides: tuple[int] = (2, 2), - batchnorm: bool = None, - dropout: float = None, - res_block_type: str = None, + nonlin: Callable = _DEFAULT_NONLIN, + n_filters: int | None = None, + conv_strides: Sequence[int] = (2, 2), + dropout: float | None = None, ) -> None: """ Bottom-up layer's method that initializes the LC modules. @@ -577,25 +638,21 @@ def _init_multiscale( which is the module responsible of merging the compressed lateral input to the main flow. - NOTE: The merge modality is set by default to "residual", meaning that the - merge layer performs concatenation on dim=1, followed by 1x1 convolution and - a Residual Gated block. + NOTE: The merge performs concatenation on dim=1, followed by 1x1 convolution and + a gated residual block. Parameters ---------- - nonlin: Callable, optional + nonlin : Callable, optional The non-linearity function used in the block. Default is `None`. - n_filters: int + n_filters : int Number of channels present through out the layers of this block. - batchnorm: bool, optional - Whether to use batchnorm layers. Default is `True`. - dropout: float, optional + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + Default is `(2, 2)`. + dropout : float, optional The dropout probability in dropout layers. If `None` dropout is not used. Default is `None`. - res_block_type: str, optional - A string specifying the structure of residual block. - Check `ResidualBlock` doscstring for more information. - Default is `None`. """ self.lowres_net = self.net if self.lowres_separate_branch: @@ -604,11 +661,8 @@ def _init_multiscale( self.lowres_merge = MergeLowRes( channels=n_filters, conv_strides=conv_strides, - merge_type="residual", nonlin=nonlin, - batchnorm=batchnorm, dropout=dropout, - res_block_type=res_block_type, multiscale_retain_spatial_dims=self.multiscale_retain_spatial_dims, multiscale_lowres_size_factor=self.multiscale_lowres_size_factor, ) @@ -620,14 +674,18 @@ def forward( Parameters ---------- - x: torch.Tensor + x : torch.Tensor The input of the `BottomUpLayer`, i.e., the input image or the output of the previous layer. - lowres_x: torch.Tensor, optional - The low-res input used for Lateral Contextualization (LC). Default is `None`. + lowres_x : torch.Tensor, optional + The low-res input used for Lateral Contextualization (LC). Default is + `None`. - NOTE: first returned tensor is used as input for the next BU layer, while the second - tensor is the bu_value passed to the top-down layer. + Returns + ------- + tuple of (torch.Tensor, torch.Tensor) + The first tensor is used as input for the next BU layer, while the second + is the `bu_value` passed to the top-down layer. """ # The input is fed through the residual downsampling block(s) primary_flow = self.net_downsized(x) @@ -640,6 +698,7 @@ def forward( return primary_flow, primary_flow if lowres_x is not None: + assert self.lowres_net is not None and self.lowres_merge is not None # First encode the low-res lateral input lowres_flow = self.lowres_net(lowres_x) # Then pass the result through the MergeLowRes layer @@ -648,24 +707,34 @@ def forward( merged = primary_flow # NOTE: Explanation of possible cases for the conditionals: - # - if both are `True` -> `merged` has the same spatial dims as the input (`x`) since - # spatial dims are retained by padding `primary_flow` in `MergeLowRes`. This is + # - if both are `True` -> `merged` has the same spatial dims as the input (`x`) + # since + # spatial dims are retained by padding `primary_flow` in `MergeLowRes`. This + # is # OK for the corresp TopDown layer, as it also retains spatial dims. - # - if both are `False` -> `merged`'s spatial dims are equal to `self.net_downsized(x)`, - # since no padding is done in `MergeLowRes` and, instead, the lowres input is cropped. - # This is OK for the corresp TopDown layer, as it also halves the spatial dims. + # - if both are `False` -> `merged`'s spatial dims are equal to + # `self.net_downsized(x)`, + # since no padding is done in `MergeLowRes` and, instead, the lowres input is + # cropped. + # This is OK for the corresp TopDown layer, as it also halves the spatial + # dims. # - if 1st is `False` and 2nd is `True` -> not a concern, it cannot happen - # (see lvae.py, line 111, intialization of `multiscale_decoder_retain_spatial_dims`). + # (see lvae.py, line 111, intialization of + # `multiscale_decoder_retain_spatial_dims`). if ( self.multiscale_retain_spatial_dims is False or self.decoder_retain_spatial_dims is True ): return merged, merged - # NOTE: if we reach here, it means that `multiscale_retain_spatial_dims` is `True`, - # but `decoder_retain_spatial_dims` is `False`, meaning that merging LC preserves - # the spatial dimensions, but at the same time we don't want to retain the spatial - # dims in the corresponding top-down layer. Therefore, we need to crop the tensor. + # NOTE: if we reach here, it means that `multiscale_retain_spatial_dims` is + # `True`, + # but `decoder_retain_spatial_dims` is `False`, meaning that merging LC + # preserves + # the spatial dimensions, but at the same time we don't want to retain the + # spatial + # dims in the corresponding top-down layer. Therefore, we need to crop the + # tensor. if self.output_expected_shape is not None: expected_shape = self.output_expected_shape else: @@ -682,115 +751,86 @@ class MergeLayer(nn.Module): """ Layer class that merges two or more input tensors. - Merges two or more (B, C, [Z], Y, X) input tensors by concatenating - them along dim=1 and passes the result through: - a) a convolutional 1x1 layer (`merge_type == "linear"`), or - b) a convolutional 1x1 layer and then a gated residual block (`merge_type == "residual"`), or - c) a convolutional 1x1 layer and then an ungated residual block (`merge_type == "residual_ungated"`). + Merges two or more (B, C, [Z], Y, X) input tensors by concatenating them along + dim=1 and passing the result through a 1x1 convolution followed by a gated + `ResidualBlock`. + + Parameters + ---------- + channels : Union[int, Iterable[int]] + The number of channels used in the convolutional blocks of this layer. + conv_strides : Sequence[int], optional + The convolution strides, used to infer the convolution dimensionality. + Default is `(2, 2)`. + nonlin : Callable, optional + The non-linearity function used in the block. Default is `nn.LeakyReLU`. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. """ def __init__( self, - merge_type: Literal["linear", "residual", "residual_ungated"], channels: Union[int, Iterable[int]], - conv_strides: tuple[int] = (2, 2), - nonlin: Callable = nn.LeakyReLU(), - batchnorm: bool = True, - dropout: Optional[float] = None, - res_block_type: Optional[str] = None, - res_block_kernel: Optional[int] = None, - conv2d_bias: Optional[bool] = True, + conv_strides: Sequence[int] = (2, 2), + nonlin: Callable = _DEFAULT_NONLIN, + dropout: float | None = None, ): """ Constructor. Parameters ---------- - merge_type: Literal["linear", "residual", "residual_ungated"] - The type of merge done in the layer. It can be chosen between "linear", - "residual", and "residual_ungated". Check the class docstring for more - information about the behaviour of different merge modalities. - channels: Union[int, Iterable[int]] + channels : Union[int, Iterable[int]] The number of channels used in the convolutional blocks of this layer. If it is an `int`: - 1st 1x1 Conv2d: in_channels=2*channels, out_channels=channels - - (Optional) ResBlock: in_channels=channels, out_channels=channels + - ResBlock: in_channels=channels, out_channels=channels If it is an Iterable (must have `len(channels)==3`): - 1st 1x1 Conv2d: in_channels=sum(channels[:-1]), out_channels=channels[-1] - - (Optional) ResBlock: in_channels=channels[-1], - out_channels=channels[-1] - conv_strides: tuple, optional + - ResBlock: in_channels=channels[-1], out_channels=channels[-1] + conv_strides : tuple, optional The strides used in the convolutions. Default is `(2, 2)`. - nonlin: Callable, optional + nonlin : Callable, optional The non-linearity function used in the block. Default is `nn.LeakyReLU`. - batchnorm: bool, optional - Whether to use batchnorm layers. Default is `True`. - dropout: float, optional + dropout : float, optional The dropout probability in dropout layers. If `None` dropout is not used. Default is `None`. - res_block_type: str, optional - A string specifying the structure of residual block. - Check `ResidualBlock` doscstring for more information. - Default is `None`. - res_block_kernel: Union[int, Iterable[int]], optional - The kernel size used in the convolutions of the residual block. - It can be either a single integer or a pair of integers defining the squared - kernel. - Default is `None`. - conv2d_bias: bool, optional - Whether to use bias term in convolutions. Default is `True`. """ super().__init__() - try: - iter(channels) - except TypeError: # it is not iterable - channels = [channels] * 3 - else: # it is iterable - if len(channels) == 1: - channels = [channels[0]] * 3 + if isinstance(channels, int): + channels_list = [channels] * 3 + else: + channels_list = list(channels) + if len(channels_list) == 1: + channels_list = [channels_list[0]] * 3 self.conv_layer: ConvType = getattr(nn, f"Conv{len(conv_strides)}d") - if merge_type == "linear": - self.layer = self.conv_layer( - sum(channels[:-1]), channels[-1], 1, bias=conv2d_bias - ) - elif merge_type == "residual": - self.layer = nn.Sequential( - self.conv_layer( - sum(channels[:-1]), channels[-1], 1, padding=0, bias=conv2d_bias - ), - ResidualGatedBlock( - conv_strides=conv_strides, - channels=channels[-1], - nonlin=nonlin, - batchnorm=batchnorm, - dropout=dropout, - block_type=res_block_type, - kernel=res_block_kernel, - conv2d_bias=conv2d_bias, - ), - ) - elif merge_type == "residual_ungated": - self.layer = nn.Sequential( - self.conv_layer( - sum(channels[:-1]), channels[-1], 1, padding=0, bias=conv2d_bias - ), - ResidualBlock( - conv_strides=conv_strides, - channels=channels[-1], - nonlin=nonlin, - batchnorm=batchnorm, - dropout=dropout, - block_type=res_block_type, - kernel=res_block_kernel, - conv2d_bias=conv2d_bias, - ), - ) + self.layer = nn.Sequential( + self.conv_layer(sum(channels_list[:-1]), channels_list[-1], 1, padding=0), + ResidualBlock( + conv_strides=conv_strides, + channels=channels_list[-1], + nonlin=nonlin, + dropout=dropout, + gated=True, + ), + ) def forward(self, *args) -> torch.Tensor: + """Concatenate the inputs along dim=1 and merge them. + + Parameters + ---------- + *args : torch.Tensor + The tensors to merge (concatenated along the channel dimension). + Returns + ------- + torch.Tensor + The merged tensor. + """ # Concatenate the input tensors along dim=1 x = torch.cat(args, dim=1) @@ -806,9 +846,27 @@ class MergeLowRes(MergeLayer): Specifically designed to merge the low-resolution patches that are used in Lateral Contextualization approach. + + Parameters + ---------- + *args : Any + Positional arguments forwarded to `MergeLayer`. + **kwargs : Any + Keyword arguments forwarded to `MergeLayer`, plus the LC-specific + `multiscale_retain_spatial_dims` and `multiscale_lowres_size_factor`. """ def __init__(self, *args, **kwargs): + """Constructor. + + Parameters + ---------- + *args : Any + Positional arguments forwarded to `MergeLayer`. + **kwargs : Any + Keyword arguments forwarded to `MergeLayer`, plus the LC-specific + `multiscale_retain_spatial_dims` and `multiscale_lowres_size_factor`. + """ self.retain_spatial_dims = kwargs.pop("multiscale_retain_spatial_dims") self.multiscale_lowres_size_factor = kwargs.pop("multiscale_lowres_size_factor") super().__init__(*args, **kwargs) @@ -818,10 +876,15 @@ def forward(self, latent: torch.Tensor, lowres: torch.Tensor) -> torch.Tensor: Parameters ---------- - latent: torch.Tensor + latent : torch.Tensor The output latent tensor from previous layer in the LVAE hierarchy. - lowres: torch.Tensor + lowres : torch.Tensor The low-res patch image to be merged to increase the context. + + Returns + ------- + torch.Tensor + The merged tensor. """ # TODO: treat (X, Y) and Z differently (e.g., line 762) if self.retain_spatial_dims: @@ -843,66 +906,6 @@ def forward(self, latent: torch.Tensor, lowres: torch.Tensor) -> torch.Tensor: return super().forward(latent, lowres) -class SkipConnectionMerger(MergeLayer): - """Specialized `MergeLayer` module, handles skip connections in the model.""" - - def __init__( - self, - nonlin: Callable, - channels: Union[int, Iterable[int]], - batchnorm: bool, - dropout: float, - res_block_type: str, - conv_strides: tuple[int] = (2, 2), - merge_type: Literal["linear", "residual", "residual_ungated"] = "residual", - conv2d_bias: bool = True, - res_block_kernel: Optional[int] = None, - ): - """ - Constructor. - - nonlin: Callable, optional - The non-linearity function used in the block. Default is `nn.LeakyReLU`. - channels: Union[int, Iterable[int]] - The number of channels used in the convolutional blocks of this layer. - If it is an `int`: - - 1st 1x1 Conv2d: in_channels=2*channels, out_channels=channels - - (Optional) ResBlock: in_channels=channels, out_channels=channels - If it is an Iterable (must have `len(channels)==3`): - - 1st 1x1 Conv2d: in_channels=sum(channels[:-1]), out_channels=channels[-1] - - (Optional) ResBlock: in_channels=channels[-1], out_channels=channels[-1] - batchnorm: bool - Whether to use batchnorm layers. - dropout: float - The dropout probability in dropout layers. If `None` dropout is not used. - res_block_type: str - A string specifying the structure of residual block. - Check `ResidualBlock` doscstring for more information. - conv_strides: tuple, optional - The strides used in the convolutions. Default is `(2, 2)`. - merge_type: Literal["linear", "residual", "residual_ungated"] - The type of merge done in the layer. It can be chosen between "linear", "residual", and "residual_ungated". - Check the class docstring for more information about the behaviour of different merge modalities. - conv2d_bias: bool, optional - Whether to use bias term in convolutions. Default is `True`. - res_block_kernel: Union[int, Iterable[int]], optional - The kernel size used in the convolutions of the residual block. - It can be either a single integer or a pair of integers defining the squared kernel. - Default is `None`. - """ - super().__init__( - conv_strides=conv_strides, - channels=channels, - nonlin=nonlin, - merge_type=merge_type, - batchnorm=batchnorm, - dropout=dropout, - res_block_type=res_block_type, - res_block_kernel=res_block_kernel, - conv2d_bias=conv2d_bias, - ) - - class TopDownLayer(nn.Module): """Top-down inference layer. @@ -923,12 +926,18 @@ class TopDownLayer(nn.Module): NOTE 2: The Top-Down layer can work in two modes: inference and prediction/generative. Depending on the particular mode, it follows distinct behaviours: - - In inference mode, parameters of q(z_i|z_i+1) are obtained from the inference path, - by merging outcomes of bottom-up and top-down passes. The exception is the top layer, - in which the parameters of q(z_L|x) are set as the output of the topmost bottom-up layer. - - On the contrary in predicition/generative mode, parameters of q(z_i|z_i+1) can be obtained - once again by merging bottom-up and top-down outputs (CONDITIONAL GENERATION), or it is - possible to directly sample from the prior p(z_i|z_i+1) (UNCONDITIONAL GENERATION). + - In inference mode, parameters of q(z_i|z_i+1) are obtained from the inference + path, + by merging outcomes of bottom-up and top-down passes. The exception is the top + layer, + in which the parameters of q(z_L|x) are set as the output of the topmost + bottom-up layer. + - On the contrary in predicition/generative mode, parameters of q(z_i|z_i+1) can + be obtained + once again by merging bottom-up and top-down outputs (CONDITIONAL GENERATION), + or it is + possible to directly sample from the prior p(z_i|z_i+1) (UNCONDITIONAL + GENERATION). NOTE 3: When doing unconditional generation, bu_value is not available. Hence the @@ -938,6 +947,43 @@ class TopDownLayer(nn.Module): If this is the top layer, at inference time, the uppermost bottom-up value is used directly as q_params, and p_params are defined in this layer (while they are usually taken from the previous layer), and can be learned. + + Parameters + ---------- + z_dim : int + The size of the latent space. + n_res_blocks : int + The number of TopDownDeterministicResBlock blocks. + n_filters : int + The number of channels present through out the layers of this block. + conv_strides : Sequence[int] + The convolution strides, used to infer the convolution dimensionality. + is_top_layer : bool, optional + Whether the current layer is at the top of the Decoder hierarchy. + Default is `False`. + upsampling_steps : int, optional + The number of upsampling steps done in this layer (typically 1). Default is + `None`. + nonlin : Callable, optional + The non-linearity function used in the block. Default is `nn.LeakyReLU`. + dropout : float, optional + The dropout probability in dropout layers. Default is `None`. + stochastic_skip : bool, optional + Whether to use a skip connection around the stochastic block. Default `False`. + learn_top_prior : bool, optional + Whether the top prior is learnable. Default is `False`. + top_prior_param_shape : Iterable[int], optional + The shape of the top-most prior parameter tensor. Default is `None`. + retain_spatial_dims : bool, optional + Whether the layer output keeps the input spatial size. Default is `False`. + vanilla_latent_hw : Iterable[int], optional + The spatial size of the latent used for prediction. Default is `None`. + input_image_shape : tuple[int, int], optional + The shape of the input image tensor. Default is `None`. + normalize_latent_factor : float, optional + A factor used to normalize the latent tensors. Default is 1.0. + stochastic_use_naive_exponential : bool, optional + Whether to use the naive (non-stable) exponential. Default is `False`. """ def __init__( @@ -945,27 +991,18 @@ def __init__( z_dim: int, n_res_blocks: int, n_filters: int, - conv_strides: tuple[int], + conv_strides: Sequence[int], is_top_layer: bool = False, upsampling_steps: Union[int, None] = None, - nonlin: Union[Callable, None] = None, - merge_type: Union[ - Literal["linear", "residual", "residual_ungated"], None - ] = None, - batchnorm: bool = True, + nonlin: Callable = _DEFAULT_NONLIN, dropout: Union[float, None] = None, stochastic_skip: bool = False, - res_block_type: Union[str, None] = None, - res_block_kernel: Union[int, None] = None, - groups: int = 1, - gated: Union[bool, None] = None, learn_top_prior: bool = False, top_prior_param_shape: Union[Iterable[int], None] = None, retain_spatial_dims: bool = False, - vanilla_latent_hw: Union[Iterable[int], None] = None, - input_image_shape: Union[tuple[int, int], None] = None, + vanilla_latent_hw: int | None = None, + input_image_shape: Sequence[int] | None = None, normalize_latent_factor: float = 1.0, - conv2d_bias: bool = True, stochastic_use_naive_exponential: bool = False, ): """ @@ -973,75 +1010,66 @@ def __init__( Parameters ---------- - z_dim: int + z_dim : int The size of the latent space. - n_res_blocks: int - The number of TopDownDeterministicResBlock blocks - n_filters: int + n_res_blocks : int + The number of TopDownDeterministicResBlock blocks. + n_filters : int The number of channels present through out the layers of this block. - conv_strides: tuple, optional + conv_strides : tuple, optional The strides used in the convolutions. Default is `(2, 2)`. - is_top_layer: bool, optional - Whether the current layer is at the top of the Decoder hierarchy. Default is `False`. - upsampling_steps: int, optional - The number of upsampling steps that has to be done in this layer (typically 1). + is_top_layer : bool, optional + Whether the current layer is at the top of the Decoder hierarchy. Default is + `False`. + upsampling_steps : int, optional + The number of upsampling steps that has to be done in this layer (typically + 1). Default is `None`. - nonlin: Callable, optional - The non-linearity function used in the block (e.g., `nn.ReLU`). Default is `None`. - merge_type: Literal["linear", "residual", "residual_ungated"], optional - The type of merge done in the layer. It can be chosen between "linear", "residual", - and "residual_ungated". Check the `MergeLayer` class docstring for more information - about the behaviour of different merging modalities. Default is `None`. - batchnorm: bool, optional - Whether to use batchnorm layers. Default is `True`. - dropout: float, optional + nonlin : Callable, optional + The non-linearity function used in the block (e.g., `nn.ReLU`). Default is + `None`. + dropout : float, optional The dropout probability in dropout layers. If `None` dropout is not used. Default is `None`. - stochastic_skip: bool, optional - Whether to use skip connections between previous top-down layer's output and this layer's stochastic output. - Stochastic skip connection allows the previous layer's output has a way to directly reach this hierarchical - level, hence facilitating the gradient flow during backpropagation. Default is `False`. - res_block_type: str, optional - A string specifying the structure of residual block. - Check `ResidualBlock` documentation for more information. - Default is `None`. - res_block_kernel: Union[int, Iterable[int]], optional - The kernel size used in the convolutions of the residual block. - It can be either a single integer or a pair of integers defining the squared kernel. - Default is `None`. - groups: int, optional - The number of groups to consider in the convolutions. Default is 1. - gated: bool, optional - Whether to use gated layer in `ResidualBlock`. Default is `None`. - learn_top_prior: + stochastic_skip : bool, optional + Whether to use skip connections between previous top-down layer's output and + this layer's stochastic output. + Stochastic skip connection allows the previous layer's output has a way to + directly reach this hierarchical + level, hence facilitating the gradient flow during backpropagation. Default + is `False`. + learn_top_prior : bool Whether to set the top prior as learnable. If this is set to `False`, in the top-most layer the prior will be N(0,1). - Otherwise, we will still have a normal distribution whose parameters will be learnt. + Otherwise, we will still have a normal distribution whose parameters will be + learnt. Default is `False`. - top_prior_param_shape: Iterable[int], optional + top_prior_param_shape : Iterable[int], optional The size of the tensor which expresses the mean and the variance of the prior for the top most layer. Default is `None`. - retain_spatial_dims: bool, optional - If `True`, the size of Encoder's latent space is kept to `input_image_shape` within the topdown layer. + retain_spatial_dims : bool, optional + If `True`, the size of Encoder's latent space is kept to `input_image_shape` + within the topdown layer. This implies that the oput spatial size equals the input spatial size. To achieve this, we centercrop the intermediate representation. Default is `False`. - vanilla_latent_hw: Iterable[int], optional - The shape of the latent tensor used for prediction (i.e., it influences the computation of restricted KL). + vanilla_latent_hw : Iterable[int], optional + The shape of the latent tensor used for prediction (i.e., it influences the + computation of restricted KL). Default is `None`. - input_image_shape: Tuple[int, int], optionalut + input_image_shape : Tuple[int, int], optionalut The shape of the input image tensor. - When `retain_spatial_dims` is set to `True`, this is used to ensure that the shape of this layer + When `retain_spatial_dims` is set to `True`, this is used to ensure that the + shape of this layer output has the same shape as the input. Default is `None`. - normalize_latent_factor: float, optional + normalize_latent_factor : float, optional A factor used to normalize the latent tensors `q_params`. - Specifically, normalization is done by dividing the latent tensor by this factor. + Specifically, normalization is done by dividing the latent tensor by this + factor. Default is 1.0. - conv2d_bias: bool, optional - Whether to use bias term is the convolutional blocks of this layer. - Default is `True`. - stochastic_use_naive_exponential: bool, optional - If `False`, in the NormalStochasticBlock2d exponentials are computed according + stochastic_use_naive_exponential : bool, optional + If `False`, in the NormalStochasticBlock2d exponentials are computed + according to the alternative definition provided by `StableExponential` class. This should improve numerical stability in the training process. Default is `False`. @@ -1053,6 +1081,7 @@ def __init__( self.stochastic_skip = stochastic_skip self.learn_top_prior = learn_top_prior self.retain_spatial_dims = retain_spatial_dims + assert input_image_shape is not None self.input_image_shape = ( input_image_shape if len(conv_strides) == 3 else input_image_shape[1:] ) @@ -1067,7 +1096,7 @@ def __init__( ) # Upsampling steps left to do in this layer - ups_left = upsampling_steps + ups_left = upsampling_steps or 0 # Define deterministic top-down block, which is a sequence of deterministic # residual blocks with (optional) upsampling. @@ -1084,13 +1113,8 @@ def __init__( conv_strides=conv_strides, nonlin=nonlin, upsample=do_resample, - batchnorm=batchnorm, dropout=dropout, - res_block_type=res_block_type, - res_block_kernel=res_block_kernel, - gated=gated, - conv2d_bias=conv2d_bias, - groups=groups, + gated=True, ) ) self.deterministic_block = nn.Sequential(*block_list) @@ -1113,34 +1137,24 @@ def __init__( self.merge = MergeLayer( channels=n_filters, conv_strides=conv_strides, - merge_type=merge_type, nonlin=nonlin, - batchnorm=batchnorm, dropout=dropout, - res_block_type=res_block_type, - res_block_kernel=res_block_kernel, - conv2d_bias=conv2d_bias, ) # Skip connection that goes around the stochastic top-down layer if stochastic_skip: - self.skip_connection_merger = SkipConnectionMerger( + self.skip_connection_merger = MergeLayer( channels=n_filters, conv_strides=conv_strides, nonlin=nonlin, - batchnorm=batchnorm, dropout=dropout, - res_block_type=res_block_type, - merge_type=merge_type, - conv2d_bias=conv2d_bias, - res_block_kernel=res_block_kernel, ) def sample_from_q( self, input_: torch.Tensor, bu_value: torch.Tensor, - var_clip_max: Optional[float] = None, + var_clip_max: float | None = None, mask: torch.Tensor = None, ) -> torch.Tensor: """ @@ -1150,21 +1164,27 @@ def sample_from_q( Parameters ---------- - input_: torch.Tensor + input_ : torch.Tensor The input tensor to the layer, which is the output of the top-down layer. - bu_value: torch.Tensor + bu_value : torch.Tensor The tensor defining the parameters /mu_q and /sigma_q computed during the bottom-up deterministic pass at the correspondent hierarchical layer. - var_clip_max: float, optional + var_clip_max : float, optional The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. Default is `None`. - mask: Union[None, torch.Tensor], optional + mask : Union[None, torch.Tensor], optional A tensor that is used to mask the sampled latent tensor. Default is `None`. + + Returns + ------- + torch.Tensor + The latent tensor sampled from q(z_i|z_{i+1}). """ if self.is_top_layer: # In top layer, we don't merge bu_value with p_params q_params = bu_value else: - # NOTE: Here the assumption is that the vampprior is only applied on the top layer. + # NOTE: Here the assumption is that the vampprior is only applied on the top + # layer. n_img_prior = None p_params = self.get_p_params(input_, n_img_prior) q_params = self.merge(bu_value, p_params) @@ -1178,8 +1198,8 @@ def sample_from_q( def get_p_params( self, - input_: torch.Tensor, - n_img_prior: int, + input_: torch.Tensor | None, + n_img_prior: int | None, ) -> torch.Tensor: """Return the parameters of the prior distribution p(z_i|z_{i+1}). @@ -1189,13 +1209,18 @@ def get_p_params( Parameters ---------- - input_: torch.Tensor - The input tensor to the layer, which is the output of the top-down layer above. - n_img_prior: int - The number of images to be generated from the unconditional prior distribution p(z_L). - """ - p_params = None + input_ : torch.Tensor or None + The input tensor to the layer, which is the output of the top-down layer + above. + n_img_prior : int or None + The number of images to be generated from the unconditional prior + distribution p(z_L). + Returns + ------- + torch.Tensor + The parameters of the prior distribution p(z_i|z_{i+1}). + """ # If top layer, define p_params as the ones of the prior p(z_L) if self.is_top_layer: p_params = self.top_prior_params @@ -1206,6 +1231,7 @@ def get_p_params( # Else the input from the layer above is p_params itself else: + assert input_ is not None p_params = input_ return p_params @@ -1227,42 +1253,48 @@ def forward( Parameters ---------- - input_: torch.Tensor, optional + input_ : torch.Tensor, optional The input tensor to the layer, which is the output of the top-down layer. Default is `None`. - skip_connection_input: torch.Tensor, optional + skip_connection_input : torch.Tensor, optional The tensor brought by the skip connection between the current and the previous top-down layer. Default is `None`. - inference_mode: bool, optional + inference_mode : bool, optional Whether the layer is in inference mode. See NOTE 2 in class description for more info. Default is `False`. - bu_value: torch.Tensor, optional + bu_value : torch.Tensor, optional The tensor defining the parameters /mu_q and /sigma_q computed during the bottom-up deterministic pass at the correspondent hierarchical layer. Default is `None`. - n_img_prior: int, optional + n_img_prior : int, optional The number of images to be generated from the unconditional prior distribution p(z_L). Default is `None`. - forced_latent: torch.Tensor, optional + forced_latent : torch.Tensor, optional A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and, hence, sampling does not happen. Default is `None`. - force_constant_output: bool, optional + force_constant_output : bool, optional Whether to copy the first sample (and rel. distrib parameters) over the whole batch. This is used when doing experiment from the prior - q is not used. Default is `False`. - mode_pred: bool, optional + mode_pred : bool, optional Whether the model is in prediction mode. Default is `False`. - use_uncond_mode: bool, optional + use_uncond_mode : bool, optional Whether to use the uncoditional distribution p(z) to sample latents in prediction mode. - var_clip_max: float + var_clip_max : float The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. + + Returns + ------- + tuple of (torch.Tensor, dict[str, torch.Tensor]) + The output tensor of the top-down layer and a dictionary of auxiliary + quantities returned by the stochastic block. """ # Check consistency of arguments inputs_none = input_ is None and skip_connection_input is None @@ -1273,6 +1305,7 @@ def forward( # Get the parameters for the latent distribution to sample from if inference_mode: # TODO What's this ? reuse Fede's code? + assert bu_value is not None if self.is_top_layer: q_params = bu_value if mode_pred is False: @@ -1295,11 +1328,12 @@ def forward( q_params = None # NOTE: Sampling is done either from q(z_i | z_{i+1}, x) or p(z_i | z_{i+1}) - # depending on the mode (hence, in practice, by checking whether q_params is None). + # depending on the mode (hence, in practice, by checking whether q_params is + # None). # Normalization of latent space parameters for stablity. # See Very deep VAEs generalize autoregressive models. - if self.normalize_latent_factor: + if self.normalize_latent_factor and q_params is not None: q_params = q_params / self.normalize_latent_factor # Sample (and process) a latent tensor in the stochastic layer @@ -1316,7 +1350,8 @@ def forward( if self.stochastic_skip and not self.is_top_layer: x = self.skip_connection_merger(x, skip_connection_input) if self.retain_spatial_dims: - # NOTE: we assume that one topdown layer will have exactly one upscaling layer. + # NOTE: we assume that one topdown layer will have exactly one upscaling + # layer. # NOTE: in case, in the Bottom-Up layer, LC retains spatial dimensions, # we have the following (see `MergeLowRes`): @@ -1332,6 +1367,7 @@ def forward( # because that's the only case in which we need to retain the shape. # Here, it must be strictly greater than half the input shape, which is # the case if and only if `x.shape == self.latent_shape`. + assert self.latent_shape is not None rescale = ( np.array((1, 2, 2)) if len(self.latent_shape) == 3 else np.array((2, 2)) ) # TODO better way? diff --git a/src/careamics/models/lvae/lvae.py b/src/careamics/models/lvae/lvae.py index 92c55e5bf..5f7c9dc04 100644 --- a/src/careamics/models/lvae/lvae.py +++ b/src/careamics/models/lvae/lvae.py @@ -5,7 +5,7 @@ and Artefact Removal, Prakash et al." """ -from collections.abc import Iterable +from collections.abc import Sequence from typing import Union import numpy as np @@ -13,6 +13,7 @@ import torch.nn as nn from careamics.models.model_utils import get_activation + from .layers import ( BottomUpDeterministicResBlock, BottomUpLayer, @@ -20,7 +21,7 @@ TopDownDeterministicResBlock, TopDownLayer, ) -from .utils import Interpolate, ModelType, crop_img_tensor +from .utils import Interpolate, crop_img_tensor class LadderVAE(nn.Module): @@ -29,8 +30,8 @@ class LadderVAE(nn.Module): Parameters ---------- - input_shape : int - The size of the input image. + input_shape : Sequence[int] + The spatial shape of the input patch, (Z, Y, X) for 3D data or (Y, X) for 2D. output_channels : int The number of output channels. multiscale_count : int @@ -64,7 +65,7 @@ class LadderVAE(nn.Module): def __init__( self, - input_shape: int, + input_shape: Sequence[int], output_channels: int, multiscale_count: int, z_dims: list[int], @@ -78,14 +79,45 @@ def __init__( encoder_blocks_per_layer: int = 1, decoder_blocks_per_layer: int = 1, ): + """Constructor. + + Parameters + ---------- + input_shape : Sequence[int] + The spatial shape of the input patch, (Z, Y, X) for 3D or (Y, X) for 2D. + output_channels : int + The number of output channels. + multiscale_count : int + The number of scales for multiscale processing. + z_dims : list[int] + The dimensions of the latent space for each layer. + encoder_n_filters : int + The number of filters in the encoder. + decoder_n_filters : int + The number of filters in the decoder. + encoder_conv_strides : list[int] + The strides for the conv layers encoder. + decoder_conv_strides : list[int] + The strides for the conv layers decoder. + encoder_dropout : float + The dropout rate for the encoder. + decoder_dropout : float + The dropout rate for the decoder. + nonlinearity : str + The nonlinearity function to use. + predict_logvar : bool + Whether to predict the log variance. + encoder_blocks_per_layer : int + The number of residual blocks per encoder layer. + decoder_blocks_per_layer : int + The number of residual blocks per decoder layer. + """ super().__init__() # ------------------------------------------------------- # Customizable attributes self.image_size = input_shape """Input image size. (Z, Y, X) or (Y, X) if the data is 2D.""" - # TODO: we need to be careful with this since used to be an int. - # the tuple of shapes used to be `self.input_shape`. self.target_ch = output_channels self.encoder_conv_strides = encoder_conv_strides self.decoder_conv_strides = decoder_conv_strides @@ -100,22 +132,12 @@ def __init__( # ------------------------------------------------------- # Model attributes -> Hardcoded - self.model_type = ModelType.LadderVae # TODO remove ! self.encoder_blocks_per_layer = encoder_blocks_per_layer self.decoder_blocks_per_layer = decoder_blocks_per_layer - self.bottomup_batchnorm = True - self.topdown_batchnorm = True - self.topdown_conv2d_bias = True - self.gated = True self.encoder_res_block_kernel = 3 - self.decoder_res_block_kernel = 3 - self.encoder_res_block_skip_padding = False - self.decoder_res_block_skip_padding = False - self.merge_type = "residual" self.no_initial_downscaling = True self.stochastic_skip = True self.learn_top_prior = True - self.res_block_type = "bacdbacd" # TODO remove ! self.mode_pred = False self._var_clip_max = 20 self._stochastic_use_naive_exponential = False @@ -201,9 +223,6 @@ def __init__( not self.no_initial_downscaling ) - # Likelihood module - # self.likelihood = self.create_likelihood_module() - # Output layer --> Project to target_ch many channels logvar_ch_needed = self.predict_logvar self.output_layer = self.parameter_net = self.decoder_conv_op( @@ -211,16 +230,8 @@ def __init__( self.target_ch * (1 + logvar_ch_needed), kernel_size=3, padding=1, - bias=self.topdown_conv2d_bias, ) - # # gradient norms. updated while training. this is also logged. - # self.grad_norm_bottom_up = 0.0 - # self.grad_norm_top_down = 0.0 - # PSNR computation on validation. - # self.label1_psnr = RunningPSNR() - # self.label2_psnr = RunningPSNR() - ### SET OF METHODS TO CREATE MODEL BLOCKS def create_first_bottom_up( self, @@ -236,10 +247,15 @@ def create_first_bottom_up( Parameters ---------- - init_stride: int + init_stride : int The stride used by the intial Conv2d block. - num_res_blocks: int, optional + num_res_blocks : int, optional The number of BottomUpDeterministicResBlocks, default is 1. + + Returns + ------- + nn.Sequential + The first bottom-up block of the Encoder. """ # From what I got from Ashesh, Z should not be touched in any case. nonlin = get_activation(self.nonlin) @@ -247,11 +263,7 @@ def create_first_bottom_up( in_channels=self.color_ch, out_channels=self.n_filters, kernel_size=self.encoder_res_block_kernel, - padding=( - 0 - if self.encoder_res_block_skip_padding - else self.encoder_res_block_kernel // 2 - ), + padding=self.encoder_res_block_kernel // 2, stride=init_stride, ) @@ -265,10 +277,7 @@ def create_first_bottom_up( c_out=self.n_filters, nonlin=nonlin, downsample=False, - batchnorm=self.bottomup_batchnorm, dropout=self.encoder_dropout, - res_block_type=self.res_block_type, - res_block_kernel=self.encoder_res_block_kernel, ) ) @@ -281,31 +290,36 @@ def create_bottom_up_layers(self, lowres_separate_branch: bool) -> nn.ModuleList that are used to generate the so-called `bu_values`. NOTE: - If `self._multiscale_count < self.n_layers`, then LC is done only in the first + If `self._multiscale_count < self.n_layers`, then LC is done only in the + first `self._multiscale_count` bottom-up layers (starting from the bottom). Parameters ---------- - lowres_separate_branch: bool + lowres_separate_branch : bool Whether the residual block(s) used for encoding the low-res input are shared (`False`) or not (`True`) with the "same-size" residual block(s) in the `BottomUpLayer`'s primary flow. + + Returns + ------- + nn.ModuleList + The stack of bottom-up layers of the Encoder. """ multiscale_lowres_size_factor = 1 nonlin = get_activation(self.nonlin) bottom_up_layers = nn.ModuleList([]) for i in range(self.n_layers): - # Whether this is the top layer - is_top = i == self.n_layers - 1 - # LC is applied only to the first (_multiscale_count - 1) bottom-up layers layer_enable_multiscale = ( self.enable_multiscale and self._multiscale_count > i + 1 ) - # This factor determines the factor by which the low-resolution tensor is larger - # N.B. Only used if layer_enable_multiscale == True, so we updated it only in that case + # This factor determines the factor by which the low-resolution tensor is + # larger + # N.B. Only used if layer_enable_multiscale == True, so we updated it only + # in that case multiscale_lowres_size_factor *= 1 + int(layer_enable_multiscale) # TODO: check correctness of this @@ -315,7 +329,8 @@ def create_bottom_up_layers(self, lowres_separate_branch: bool) -> nn.ModuleList output_expected_shape = None # Add bottom-up deterministic layer at level i. - # It's a sequence of residual blocks (BottomUpDeterministicResBlock), possibly with downsampling between them. + # It's a sequence of residual blocks (BottomUpDeterministicResBlock), + # possibly with downsampling between them. bottom_up_layers.append( BottomUpLayer( n_res_blocks=self.encoder_blocks_per_layer, @@ -323,17 +338,22 @@ def create_bottom_up_layers(self, lowres_separate_branch: bool) -> nn.ModuleList downsampling_steps=self.downsample[i], nonlin=nonlin, conv_strides=self.encoder_conv_strides, - batchnorm=self.bottomup_batchnorm, dropout=self.encoder_dropout, - res_block_type=self.res_block_type, - res_block_kernel=self.encoder_res_block_kernel, - gated=self.gated, lowres_separate_branch=lowres_separate_branch, - enable_multiscale=self.enable_multiscale, # TODO: shouldn't the arg be `layer_enable_multiscale` here? + # NOTE: the global `enable_multiscale` flag is passed to every layer + # (not the per-layer `layer_enable_multiscale`). As a result, layers + # above index `_multiscale_count - 1` build + # `lowres_net`/`lowres_merge` modules that are never exercised + # (they only ever receive `lowres_x=None` and return the primary + # flow unchanged, so the output is identical either way). Passing + # `layer_enable_multiscale` would drop those unused parameters but + # change the `state_dict` and thus break existing checkpoints, so it + # is left as-is intentionally. + enable_multiscale=self.enable_multiscale, multiscale_retain_spatial_dims=self.multiscale_retain_spatial_dims, - multiscale_lowres_size_factor=multiscale_lowres_size_factor, + multiscale_lowres_size_factor=multiscale_lowres_size_factor, # type: ignore[arg-type] decoder_retain_spatial_dims=self.multiscale_decoder_retain_spatial_dims, - output_expected_shape=output_expected_shape, + output_expected_shape=output_expected_shape, # type: ignore[arg-type] ) ) @@ -343,9 +363,12 @@ def create_top_down_layers(self) -> nn.ModuleList: """ Method creates the stack of top-down layers of the Decoder. - In these layer the `bu`_values` from the Encoder are merged with the `p_params` from the previous layer - of the Decoder to get `q_params`. Then, a stochastic layer generates a sample from the latent distribution - with parameters `q_params`. Finally, this sample is fed through a TopDownDeterministicResBlock to + In these layer the `bu`_values` from the Encoder are merged with the `p_params` + from the previous layer + of the Decoder to get `q_params`. Then, a stochastic layer generates a sample + from the latent distribution + with parameters `q_params`. Finally, this sample is fed through a + TopDownDeterministicResBlock to compute the `p_params` for the layer below. NOTE 1: @@ -361,6 +384,10 @@ def create_top_down_layers(self) -> nn.ModuleList: When doing unconditional generation, bu_value is not available. Hence the merge layer is not used, and z is sampled directly from p_params. + Returns + ------- + nn.ModuleList + The stack of top-down layers of the Decoder. """ top_down_layers = nn.ModuleList([]) nonlin = get_activation(self.nonlin) @@ -385,20 +412,14 @@ def create_top_down_layers(self) -> nn.ModuleList: conv_strides=self.decoder_conv_strides, upsampling_steps=self.downsample[i], nonlin=nonlin, - merge_type=self.merge_type, - batchnorm=self.topdown_batchnorm, dropout=self.decoder_dropout, stochastic_skip=self.stochastic_skip, learn_top_prior=self.learn_top_prior, top_prior_param_shape=self.get_top_prior_param_shape(), - res_block_type=self.res_block_type, - res_block_kernel=self.decoder_res_block_kernel, - gated=self.gated, vanilla_latent_hw=self.get_latent_spatial_size(i), retain_spatial_dims=self.multiscale_decoder_retain_spatial_dims, input_image_shape=self.image_size, normalize_latent_factor=normalize_latent_factor, - conv2d_bias=self.topdown_conv2d_bias, stochastic_use_naive_exponential=self._stochastic_use_naive_exponential, ) ) @@ -407,39 +428,41 @@ def create_top_down_layers(self) -> nn.ModuleList: def create_final_topdown_layer(self, upsample: bool) -> nn.Sequential: """Create the final top-down layer of the Decoder. - NOTE: In this layer, (optional) upsampling is performed by bilinear interpolation + NOTE: In this layer, (optional) upsampling is performed by bilinear + interpolation instead of transposed convolution (like in other TD layers). Parameters ---------- - upsample: bool + upsample : bool Whether to upsample the input of the final top-down layer by bilinear interpolation with `scale_factor=2`. + + Returns + ------- + nn.Sequential + The final top-down layer of the Decoder. """ # Final top-down layer - modules = list() + modules = [] if upsample: modules.append(Interpolate(scale=2)) - for i in range(self.decoder_blocks_per_layer): + for _i in range(self.decoder_blocks_per_layer): modules.append( TopDownDeterministicResBlock( c_in=self.n_filters, c_out=self.n_filters, nonlin=get_activation(self.nonlin), conv_strides=self.decoder_conv_strides, - batchnorm=self.topdown_batchnorm, dropout=self.decoder_dropout, - res_block_type=self.res_block_type, - res_block_kernel=self.decoder_res_block_kernel, - gated=self.gated, - conv2d_bias=self.topdown_conv2d_bias, + gated=True, ) ) return nn.Sequential(*modules) - def _init_multires(self, config=None) -> nn.ModuleList: + def _init_multires(self) -> None: """ Method defines the input block/branch to encode/compress low-res lateral inputs. @@ -472,12 +495,6 @@ def _init_multires(self, config=None) -> nn.ModuleList: self._multiscale_count <= 1 or self._multiscale_count <= 1 + self.n_layers ), msg # TODO how ? - msg = ( - "Multiscale approach only supports monocrome images. " - f"Found instead color_ch={self.color_ch}." - ) - # assert self._multiscale_count == 1 or self.color_ch == 1, msg - lowres_first_bottom_ups = [] for _ in range(1, self._multiscale_count): first_bottom_up = nn.Sequential( @@ -495,9 +512,7 @@ def _init_multires(self, config=None) -> nn.ModuleList: conv_strides=self.encoder_conv_strides, nonlin=nonlin, downsample=False, - batchnorm=self.bottomup_batchnorm, dropout=self.encoder_dropout, - res_block_type=self.res_block_type, ), ) lowres_first_bottom_ups.append(first_bottom_up) @@ -510,22 +525,6 @@ def _init_multires(self, config=None) -> nn.ModuleList: ### SET OF FORWARD-LIKE METHODS def bottomup_pass(self, inp: torch.Tensor) -> list[torch.Tensor]: - """Wrapper of _bottomup_pass().""" - # TODO Remove wrapper - return self._bottomup_pass( - inp, - self.first_bottom_up, - self.lowres_first_bottom_ups, - self.bottom_up_layers, - ) - - def _bottomup_pass( - self, - inp: torch.Tensor, - first_bottom_up: nn.Sequential, - lowres_first_bottom_ups: nn.ModuleList, - bottom_up_layers: nn.ModuleList, - ) -> list[torch.Tensor]: """ Method defines the forward pass through the LVAE Encoder, the so-called. @@ -533,22 +532,22 @@ def _bottomup_pass( Parameters ---------- - inp: torch.Tensor - The input tensor to the bottom-up pass of shape (B, 1+n_LC, H, W), where n_LC + inp : torch.Tensor + The input tensor to the bottom-up pass of shape (B, 1+n_LC, H, W), where + n_LC is the number of lateral low-res inputs used in the LC approach. In particular, the first channel corresponds to the input patch, while the remaining ones are associated to the lateral low-res inputs. - first_bottom_up: nn.Sequential - The module defining the first bottom-up layer of the Encoder. - lowres_first_bottom_ups: nn.ModuleList - The list of modules defining Lateral Contextualization. - bottom_up_layers: nn.ModuleList - The list of modules defining the stack of bottom-up layers of the Encoder. + + Returns + ------- + list[torch.Tensor] + The `bu_values`, one deterministic tensor per bottom-up layer. """ if self._multiscale_count > 1: - x = first_bottom_up(inp[:, :1]) + x = self.first_bottom_up(inp[:, :1]) else: - x = first_bottom_up(inp) + x = self.first_bottom_up(inp) # Loop from bottom to top layer, store all deterministic nodes we # need for the top-down pass in bu_values list @@ -556,8 +555,9 @@ def _bottomup_pass( for i in range(self.n_layers): lowres_x = None if self._multiscale_count > 1 and i + 1 < inp.shape[1]: - lowres_x = lowres_first_bottom_ups[i](inp[:, i + 1 : i + 2]) - x, bu_value = bottom_up_layers[i](x, lowres_x=lowres_x) + assert self.lowres_first_bottom_ups is not None + lowres_x = self.lowres_first_bottom_ups[i](inp[:, i + 1 : i + 2]) + x, bu_value = self.bottom_up_layers[i](x, lowres_x=lowres_x) bu_values.append(bu_value) return bu_values @@ -566,7 +566,7 @@ def topdown_pass( self, bu_values: Union[torch.Tensor, None] = None, n_img_prior: Union[torch.Tensor, None] = None, - constant_layers: Union[Iterable[int], None] = None, + constant_layers: Union[Sequence[int], None] = None, forced_latent: Union[list[torch.Tensor], None] = None, top_down_layers: Union[nn.ModuleList, None] = None, final_top_down_layer: Union[nn.Sequential, None] = None, @@ -578,26 +578,32 @@ def topdown_pass( Parameters ---------- - bu_values: torch.Tensor, optional + bu_values : torch.Tensor, optional Output of the bottom-up pass. It will have values from multiple layers of the ladder. - n_img_prior: optional + n_img_prior : optional When `bu_values` is `None`, `n_img_prior` indicates the number of images to generate from the prior (so bottom-up pass is not used at all here). - constant_layers: Iterable[int], optional + constant_layers : Iterable[int], optional A sequence of indexes associated to the layers in which a single instance's z is copied over the entire batch (bottom-up path is not used, so only prior is used here). Set to `None` to avoid this behaviour. - forced_latent: list[torch.Tensor], optional + forced_latent : list[torch.Tensor], optional A list of tensors that are used as fixed latent variables (hence, sampling doesn't take place in this case). - top_down_layers: nn.ModuleList, optional + top_down_layers : nn.ModuleList, optional A list of top-down layers to use in the top-down pass. If `None`, the method uses the default layers defined in the constructor. - final_top_down_layer: nn.Sequential, optional + final_top_down_layer : nn.Sequential, optional The last top-down layer of the top-down pass. If `None`, the method uses the default layers defined in the constructor. + + Returns + ------- + tuple of (torch.Tensor, dict[str, torch.Tensor]) + The output tensor of the top-down pass and a dictionary of auxiliary + quantities (sampled latents, KL terms, distribution parameters, etc.). """ if top_down_layers is None: top_down_layers = self.top_down_layers @@ -649,7 +655,7 @@ def topdown_pass( for i in reversed(range(self.n_layers)): # If available, get deterministic node from bottom-up inference try: - bu_value = bu_values[i] + bu_value = bu_values[i] # type: ignore[index] except TypeError: bu_value = None @@ -681,10 +687,6 @@ def topdown_pass( kl_channelwise[i] = aux["kl_channelwise"] debug_qvar_max[i] = aux["qvar_max"] - # if self.mode_pred is False: - # logprob_p += aux['logprob_p'].mean() # mean over batch - # else: - # logprob_p = None # Final top-down layer out = final_top_down_layer(out) @@ -695,8 +697,7 @@ def topdown_pass( "kl": kl, # list of tensors with shape (batch, ) "kl_restricted": kl_restricted, # list of tensors with shape (batch, ) "kl_spatial": kl_spatial, # list of tensors w shape (batch, h[i], w[i]) - "kl_channelwise": kl_channelwise, # list of tensors with shape (batch, ch[i]) - # 'logprob_p': logprob_p, # scalar, mean over batch + "kl_channelwise": kl_channelwise, # list of tensors, shape (batch, ch[i]) "q_mu": q_mu, "q_lv": q_lv, "debug_qvar_max": debug_qvar_max, @@ -709,8 +710,13 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor Parameters ---------- - x: torch.Tensor + x : torch.Tensor The input tensor of shape (B, C, H, W). + + Returns + ------- + tuple of (torch.Tensor, dict[str, torch.Tensor]) + The output tensor and a dictionary of auxiliary top-down quantities. """ img_size = x.size()[2:] @@ -718,6 +724,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor bu_values = self.bottomup_pass(x) if self._squish3d: + assert self._3D_squisher is not None bu_values = [ torch.mean(self._3D_squisher[k](bu_value), dim=2) for k, bu_value in enumerate(bu_values) @@ -736,20 +743,26 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor ### SET OF GETTERS def get_padded_size(self, size): - """ - Returns the smallest size (H, W) of the image with actual size given - as input, such that H and W are powers of 2. - :param size: input size, tuple either (N, C, H, W) or (H, W) - :return: 2-tuple (H, W) + """Return the smallest (H, W) padded size that are powers of 2. + + Parameters + ---------- + size : Sequence[int] + The input size, either (N, C, H, W) or (H, W). + + Returns + ------- + list[int] + The padded (H, W) size. """ # Make size argument into (heigth, width) - # assert len(size) in [2, 4, 5] # TODO commented out cuz it's weird # We're only interested in the Y,X dimensions size = size[-2:] if self.multiscale_decoder_retain_spatial_dims is True: # In this case, we can go much more deeper and so this is not required - # (in the way it is. ;). More work would be needed if this was to be correctly implemented ) + # (in the way it is. ;). More work would be needed if this was to be + # correctly implemented ) return list(size) # Overall downscale factor from input to top layer (power of 2) @@ -761,7 +774,18 @@ def get_padded_size(self, size): return padded_size def get_latent_spatial_size(self, level_idx: int): - """Level_idx: 0 is the bottommost layer, the highest resolution one.""" + """Return the spatial size of the latent tensor at the given level. + + Parameters + ---------- + level_idx : int + The hierarchy level; 0 is the bottommost (highest-resolution) layer. + + Returns + ------- + int + The spatial size (height, equal to width) of the latent at that level. + """ actual_downsampling = level_idx + 1 dwnsc = 2**actual_downsampling sz = self.get_padded_size(self.image_size) @@ -770,34 +794,52 @@ def get_latent_spatial_size(self, level_idx: int): assert h == w return h - def get_top_prior_param_shape(self, n_imgs: int = 1): + def get_top_prior_param_shape(self, n_imgs: int = 1) -> tuple[int, ...]: + """Compute the shape of the top-layer prior parameter tensor. + Parameters + ---------- + n_imgs : int, optional + The number of images (batch dimension) in the prior tensor. Default is 1. + + Returns + ------- + tuple[int, ...] + The shape of the top prior parameter tensor (mean and log-variance stacked). + """ # Compute the total downscaling performed in the Encoder if self.multiscale_decoder_retain_spatial_dims is False: dwnsc = self.overall_downscale_factor else: - # LC allow the encoder latents to keep the same (H, W) size at different levels + # LC allow the encoder latents to keep the same (H, W) size at different + # levels actual_downsampling = self.n_layers + 1 - self._multiscale_count dwnsc = 2**actual_downsampling h = self.image_size[-2] // dwnsc w = self.image_size[-1] // dwnsc mu_logvar = self.z_dims[-1] * 2 # mu and logvar - top_layer_shape = (n_imgs, mu_logvar, h, w) + top_layer_shape: tuple[int, ...] = (n_imgs, mu_logvar, h, w) # TODO refactor! if self._model_3D_depth > 1 and self._decoder_mode_3D is True: # TODO check if model_3D_depth is needed ? top_layer_shape = (n_imgs, mu_logvar, self._model_3D_depth, h, w) return top_layer_shape - def reset_for_inference(self, tile_size: tuple[int, int] | None = None): - """Should be called if we want to predict for a different input/output size.""" + def reset_for_inference(self, tile_size: Sequence[int] | None = None): + """Reconfigure the model to predict for a different input/output size. + + Parameters + ---------- + tile_size : Sequence[int], optional + The new spatial tile size. If `None`, the training `image_size` is kept. + """ self.mode_pred = True if tile_size is None: tile_size = self.image_size self.image_size = tile_size for i in range(self.n_layers): - self.bottom_up_layers[i].output_expected_shape = ( + self.bottom_up_layers[i].output_expected_shape = tuple( ts // 2 ** (i + 1) for ts in tile_size ) self.top_down_layers[i].latent_shape = tile_size diff --git a/src/careamics/models/lvae/stochastic.py b/src/careamics/models/lvae/stochastic.py index 25a6be027..ad6e03343 100644 --- a/src/careamics/models/lvae/stochastic.py +++ b/src/careamics/models/lvae/stochastic.py @@ -1,8 +1,6 @@ -"""Script containing the common basic blocks (nn.Module) -reused by the LadderVAE architecture. -""" +"""Stochastic latent block (nn.Module) used by the LadderVAE architecture.""" -from typing import Dict, Tuple, Union +from typing import Union import torch import torch.nn as nn @@ -33,10 +31,33 @@ class NormalStochasticBlock(nn.Module): If parameters for q are not given, sampling is done from p(z). NOTE 2: - The restricted KL divergence is obtained by first computing the element-wise KL divergence - (i.e., the KL computed for each element of the latent tensors). Then, the restricted version - is computed by summing over the channels and the spatial dimensions associated only to the + The restricted KL divergence is obtained by first computing the element-wise KL + divergence + (i.e., the KL computed for each element of the latent tensors). Then, the + restricted version + is computed by summing over the channels and the spatial dimensions associated + only to the portion of the latent tensor that is used for prediction. + + Parameters + ---------- + c_in : int + The number of channels of the input tensor. + c_vars : int + The number of channels of the latent space tensor. + c_out : int + The number of output channels of the stochastic layer (differs from the latent + z). + conv_dims : int, optional + The number of dimensions of the convolutional layers (2D or 3D). Default is 2. + kernel : int, optional + The size of the kernel used in the convolutional layers. Default is 3. + transform_p_params : bool, optional + Whether to apply a transformation to the `p_params` tensor. Default is `True`. + vanilla_latent_hw : int, optional + The spatial size of the latent tensor used for prediction. Default is `None`. + use_naive_exponential : bool, optional + Whether to use the naive (non-stable) exponential. Default is `False`. """ def __init__( @@ -47,36 +68,40 @@ def __init__( conv_dims: int = 2, kernel: int = 3, transform_p_params: bool = True, - vanilla_latent_hw: int = None, + vanilla_latent_hw: int | None = None, use_naive_exponential: bool = False, ): - """ + """Constructor. + Parameters ---------- - c_in: int + c_in : int The number of channels of the input tensor. - c_vars: int + c_vars : int The number of channels of the latent space tensor. - c_out: int + c_out : int The output of the stochastic layer. Note that this is different from the sampled latent z. - conv_dims: int, optional + conv_dims : int, optional The number of dimensions of the convolutional layers (2D or 3D). Default is 2. - kernel: int, optional + kernel : int, optional The size of the kernel used in convolutional layers. Default is 3. - transform_p_params: bool, optional + transform_p_params : bool, optional Whether a transformation should be applied to the `p_params` tensor. The transformation consists in a 2D convolution ()`conv_in_p()`) that maps the input to a larger number of channels. Default is `True`. - vanilla_latent_hw: int, optional - The shape of the latent tensor used for prediction (i.e., it influences the computation of restricted KL). + vanilla_latent_hw : int, optional + The shape of the latent tensor used for prediction (i.e., it influences the + computation of restricted KL). Default is `None`. - use_naive_exponential: bool, optional - If `False`, exponentials are computed according to the alternative definition - provided by `StableExponential` class. This should improve numerical stability + use_naive_exponential : bool, optional + If `False`, exponentials are computed according to the alternative + definition + provided by `StableExponential` class. This should improve numerical + stability in the training process. Default is `False`. """ super().__init__() @@ -110,19 +135,27 @@ def get_z( - Sampled from the (Gaussian) latent distribution. - Taken as a pre-defined forced latent. - Taken as the mode (mean) of the latent distribution. - - In prediction mode (`mode_pred==True`), can be either sample or taken as the distribution mode. + - In prediction mode (`mode_pred==True`), can be either sample or taken as + the distribution mode. Parameters ---------- - sampling_distrib: torch.distributions.normal.Normal + sampling_distrib : torch.distributions.normal.Normal The Gaussian distribution from which latent tensor is sampled. - forced_latent: torch.Tensor - A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent tensor and, + forced_latent : torch.Tensor + A pre-defined latent tensor. If it is not `None`, than it is used as the + actual latent tensor and, hence, sampling does not happen. - mode_pred: bool + mode_pred : bool Whether the model is prediction mode. - use_uncond_mode: bool - Whether to use the uncoditional distribution p(z) to sample latents in prediction mode. + use_uncond_mode : bool + Whether to use the uncoditional distribution p(z) to sample latents in + prediction mode. + + Returns + ------- + torch.Tensor + The sampled (or forced) latent tensor. """ if forced_latent is None: if mode_pred: @@ -137,20 +170,25 @@ def get_z( return z def sample_from_q( - self, q_params: torch.Tensor, var_clip_max: float + self, q_params: torch.Tensor, var_clip_max: float | None ) -> torch.Tensor: - """ - Given an input parameter tensor defining q(z), - it processes it by calling `process_q_params()` method and - sample a latent tensor from the resulting distribution. + """Sample a latent tensor from the inference distribution q(z). + + The input parameter tensor defining q(z) is processed by `process_q_params()` + and a latent tensor is sampled from the resulting distribution. Parameters ---------- - q_params: torch.Tensor + q_params : torch.Tensor The input tensor to be processed. - var_clip_max: float + var_clip_max : float or None The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. + + Returns + ------- + torch.Tensor + The latent tensor sampled from q(z). """ _, _, q = self.process_q_params(q_params, var_clip_max) return q.rsample() @@ -163,31 +201,42 @@ def compute_kl_metrics( q_params: torch.Tensor, mode_pred: bool, z: torch.Tensor, - ) -> Dict[str, torch.Tensor]: - """ - Compute the (Monte Carlo estimated) KL and extract composed versions of the metric. + ) -> dict[str, torch.Tensor]: + """Compute the (Monte Carlo estimated) KL and its composed versions. + Specifically, the different versions of the KL loss terms are: - - `kl_elementwise`: KL term for each single element of the latent tensor [Shape: (batch, ch, h, w)]. - - `kl_samplewise`: KL term associated to each sample in the batch [Shape: (batch, )]. - - `kl_samplewise_restricted`: KL term only associated to the portion of the latent tensor that is - used for prediction and summed over channel and spatial dimensions [Shape: (batch, )]. - - `kl_channelwise`: KL term associated to each sample and each channel [Shape: (batch, ch, )]. - - `kl_spatial`: KL term summed over the channels, i.e., retaining the spatial dimensions [Shape: (batch, h, w)] + - `kl_elementwise`: KL term for each single element of the latent tensor + [Shape: (batch, ch, h, w)]. + - `kl_samplewise`: KL term associated to each sample in the batch [Shape: + (batch, )]. + - `kl_samplewise_restricted`: KL term only associated to the portion of the + latent tensor that is + used for prediction and summed over channel and spatial dimensions [Shape: + (batch, )]. + - `kl_channelwise`: KL term associated to each sample and each channel + [Shape: (batch, ch, )]. + - `kl_spatial`: KL term summed over the channels, i.e., retaining the + spatial dimensions [Shape: (batch, h, w)]. Parameters ---------- - p: torch.distributions.normal.Normal + p : torch.distributions.normal.Normal The prior generative distribution p(z_i|z_{i+1}) (or p(z_L)). - p_params: torch.Tensor + p_params : torch.Tensor The parameters of the prior generative distribution. - q: torch.distributions.normal.Normal + q : torch.distributions.normal.Normal The inference distribution q(z_i|z_{i+1}) (or q(z_L|x)). - q_params: torch.Tensor + q_params : torch.Tensor The parameters of the inference distribution. - mode_pred: bool + mode_pred : bool Whether the model is in prediction mode. - z: torch.Tensor + z : torch.Tensor The sampled latent tensor. + + Returns + ------- + dict[str, torch.Tensor] + The dictionary of KL metrics (see the method summary for the keys). """ kl_samplewise_restricted = None if mode_pred is False: # if not predicting @@ -197,7 +246,8 @@ def compute_kl_metrics( kl_samplewise = kl_elementwise.sum(all_dims[1:]) kl_channelwise = kl_elementwise.sum(all_dims[2:]) - # compute KL only on the portion of the latent space that is used for prediction. + # compute KL only on the portion of the latent space that is used for + # prediction. pad = (kl_elementwise.shape[-1] - self._vanilla_latent_hw) // 2 if pad > 0: tmp = kl_elementwise[..., pad:-pad, pad:-pad] @@ -221,23 +271,29 @@ def compute_kl_metrics( return kl_dict def process_p_params( - self, p_params: torch.Tensor, var_clip_max: float - ) -> Tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]: - """Process the input parameters to get the prior distribution p(z_i|z_{i+1}) (or p(z_L)). + self, p_params: torch.Tensor, var_clip_max: float | None + ) -> tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]: + """Process the input parameters to get the prior distribution p(z_i|z_{i+1}). Processing consists in: - - (optionally) 2D convolution on the input tensor to increase number of channels. + - (optionally) 2D convolution on the input tensor to increase number of + channels. - split the resulting tensor into two chunks, the mean and the log-variance. - (optionally) clip the log-variance to an upper threshold. - define the normal distribution p(z) given the parameter tensors above. Parameters ---------- - p_params: torch.Tensor + p_params : torch.Tensor The input tensor to be processed. - var_clip_max: float + var_clip_max : float or None The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. + + Returns + ------- + tuple of (torch.Tensor, torch.Tensor, torch.distributions.normal.Normal) + The prior mean wrapper, log-variance wrapper, and the distribution p(z). """ if self.transform_p_params: p_params = self.conv_in_p(p_params) @@ -255,25 +311,35 @@ def process_p_params( return p_mu, p_lv, p def process_q_params( - self, q_params: torch.Tensor, var_clip_max: float, allow_oddsizes: bool = False - ) -> Tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]: - """ - Process the input parameters to get the inference distribution q(z_i|z_{i+1}) (or q(z|x)). + self, + q_params: torch.Tensor, + var_clip_max: float | None, + allow_oddsizes: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, torch.distributions.normal.Normal]: + """Process the input parameters into the inference distribution q(z). Processing consists in: - convolution on the input tensor to double the number of channels. - split the resulting tensor into 2 chunks, respectively mean and log-var. - (optionally) clip the log-variance to an upper threshold. - - (optionally) crop the resulting tensors to ensure that the last spatial dimension is even. + - (optionally) crop the resulting tensors to keep the last spatial dim even. - define the normal distribution q(z) given the parameter tensors above. Parameters ---------- - p_params: torch.Tensor + q_params : torch.Tensor The input tensor to be processed. - var_clip_max: float + var_clip_max : float or None The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. + allow_oddsizes : bool, optional + Whether to allow an odd last spatial dimension (skip the centercrop). + Default is `False`. + + Returns + ------- + tuple of (torch.Tensor, torch.Tensor, torch.distributions.normal.Normal) + The inference mean wrapper, log-variance wrapper, and the distribution q(z). """ q_params = self.conv_in_q(q_params) @@ -299,30 +365,42 @@ def forward( mode_pred: bool = False, use_uncond_mode: bool = False, var_clip_max: Union[float, None] = None, - ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: - """ + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Run the stochastic block: sample a latent and compute KL metrics. + Parameters ---------- - p_params: torch.Tensor - The output tensor of the top-down layer above (i.e., mu_{p,i+1}, sigma_{p,i+1}). - q_params: torch.Tensor, optional - The tensor resulting from merging the bu_value tensor at the same hierarchical level + p_params : torch.Tensor + The output tensor of the top-down layer above (i.e., mu_{p,i+1}, + sigma_{p,i+1}). + q_params : torch.Tensor, optional + The tensor resulting from merging the bu_value tensor at the same + hierarchical level from the bottom-up pass and the `p_params` tensor. Default is `None`. - forced_latent: torch.Tensor, optional - A pre-defined latent tensor. If it is not `None`, than it is used as the actual latent + forced_latent : torch.Tensor, optional + A pre-defined latent tensor. If it is not `None`, than it is used as the + actual latent tensor and, hence, sampling does not happen. Default is `None`. - force_constant_output: bool, optional - Whether to copy the first sample (and rel. distrib parameters) over the whole batch. + force_constant_output : bool, optional + Whether to copy the first sample (and rel. distrib parameters) over the + whole batch. This is used when doing experiment from the prior - q is not used. Default is `False`. - mode_pred: bool, optional + mode_pred : bool, optional Whether the model is in prediction mode. Default is `False`. - use_uncond_mode: bool, optional - Whether to use the uncoditional distribution p(z) to sample latents in prediction mode. + use_uncond_mode : bool, optional + Whether to use the uncoditional distribution p(z) to sample latents in + prediction mode. Default is `False`. - var_clip_max: float, optional + var_clip_max : float, optional The maximum value reachable by the log-variance of the latent distribution. Values exceeding this threshold are clipped. Default is `None`. + + Returns + ------- + tuple of (torch.Tensor, dict[str, torch.Tensor]) + The output tensor of the stochastic block and a dictionary of auxiliary + quantities (sampled latent, distribution parameters, KL metrics, etc.). """ debug_qvar_max = 0 diff --git a/src/careamics/models/lvae/utils.py b/src/careamics/models/lvae/utils.py index 2698dbf5a..690784b96 100644 --- a/src/careamics/models/lvae/utils.py +++ b/src/careamics/models/lvae/utils.py @@ -1,8 +1,7 @@ -""" -Script for utility functions needed by the LVAE model. -""" +"""Script for utility functions needed by the LVAE model.""" -from typing import Literal, Sequence +from collections.abc import Sequence +from typing import Literal import numpy as np import torch @@ -11,45 +10,86 @@ from torch.distributions.normal import Normal -def torch_nanmean(inp): - return torch.mean(inp[~inp.isnan()]) +def torch_nanmean(inp: torch.Tensor) -> torch.Tensor: + """Compute the mean of a tensor ignoring NaN values. + Parameters + ---------- + inp : torch.Tensor + Input tensor, possibly containing NaN values. -def power_of_2(self, x): - assert isinstance(x, int) - if x == 1: - return True - if x == 0: - # happens with validation - return False - if x % 2 == 1: - return False - return self.power_of_2(x // 2) + Returns + ------- + torch.Tensor + The mean of the non-NaN elements of the input. + """ + return torch.mean(inp[~inp.isnan()]) class Enum: + """Lightweight enum-like base class backed by class attributes.""" + @classmethod - def name(cls, enum_type): + def name(cls, enum_type: int) -> str | None: + """Return the attribute name matching the given value. + + Parameters + ---------- + enum_type : int + The value to look up. + + Returns + ------- + str or None + The name of the attribute holding the given value, or `None` if not found. + """ for key, value in cls.__dict__.items(): if enum_type == value: return key + return None @classmethod - def contains(cls, enum_type): - for key, value in cls.__dict__.items(): + def contains(cls, enum_type: int) -> bool: + """Return whether the given value is defined in the enum. + + Parameters + ---------- + enum_type : int + The value to look up. + + Returns + ------- + bool + Whether the value is defined in the enum. + """ + for _key, value in cls.__dict__.items(): if enum_type == value: return True return False @classmethod - def from_name(cls, enum_type_str): + def from_name(cls, enum_type_str: str) -> int: + """Return the value for the given attribute name. + + Parameters + ---------- + enum_type_str : str + The name of the attribute to look up. + + Returns + ------- + int + The value held by the attribute. + """ for key, value in cls.__dict__.items(): if key == enum_type_str: return value - assert f"{cls.__name__}:{enum_type_str} doesnot exist." + raise ValueError(f"{cls.__name__}:{enum_type_str} does not exist.") class LossType(Enum): + """Enumeration of the loss types supported by the LVAE training code.""" + Elbo = 0 ElboWithCritic = 1 ElboMixedReconstruction = 2 @@ -61,60 +101,24 @@ class LossType(Enum): DenoiSplitMuSplit = 8 -class ModelType(Enum): - LadderVae = 3 - LadderVaeTwinDecoder = 4 - LadderVAECritic = 5 - # Separate vampprior: two optimizers - LadderVaeSepVampprior = 6 - # one encoder for mixed input, two for separate inputs. - LadderVaeSepEncoder = 7 - LadderVAEMultiTarget = 8 - LadderVaeSepEncoderSingleOptim = 9 - UNet = 10 - BraveNet = 11 - LadderVaeStitch = 12 - LadderVaeSemiSupervised = 13 - LadderVaeStitch2Stage = 14 # Note that previously trained models will have issue. - # since earlier, LadderVaeStitch2Stage = 13, LadderVaeSemiSupervised = 14 - LadderVaeMixedRecons = 15 - LadderVaeCL = 16 - LadderVaeTwoDataSet = ( - 17 # on one subdset, apply disentanglement, on other apply reconstruction - ) - LadderVaeTwoDatasetMultiBranch = 18 - LadderVaeTwoDatasetMultiOptim = 19 - LVaeDeepEncoderIntensityAug = 20 - AutoRegresiveLadderVAE = 21 - LadderVAEInterleavedOptimization = 22 - Denoiser = 23 - DenoiserSplitter = 24 - SplitterDenoiser = 25 - LadderVAERestrictedReconstruction = 26 - LadderVAETwoDataSetRestRecon = 27 - LadderVAETwoDataSetFinetuning = 28 - - def _pad_crop_img( x: torch.Tensor, size: Sequence[int], mode: Literal["crop", "pad"] ) -> torch.Tensor: - """Pads or crops a tensor. - - Pads or crops a tensor of shape (B, C, [Z], Y, X) to new shape. - - Parameters: - ----------- - x: torch.Tensor - Input image of shape (B, C, [Z], Y, X) - size: Sequence[int] - Desired size ([Z*], Y*, X*) - mode: Literal["crop", "pad"] - Mode, either 'pad' or 'crop' - - Returns: - -------- - torch.Tensor: - The padded or cropped tensor + """Pad or crop a tensor of shape (B, C, [Z], Y, X) to a new spatial shape. + + Parameters + ---------- + x : torch.Tensor + Input image of shape (B, C, [Z], Y, X). + size : Sequence[int] + Desired spatial size ([Z*], Y*, X*). + mode : Literal["crop", "pad"] + Whether to 'pad' or 'crop' the input. + + Returns + ------- + torch.Tensor + The padded or cropped tensor. """ # TODO: Support cropping/padding on selected dimensions assert (x.dim() == 4 and len(size) == 2) or (x.dim() == 5 and len(size) == 3) @@ -130,7 +134,7 @@ def _pad_crop_img( if cond: raise ValueError(f"Trying to {mode} from size {x_size} to size {size}") - diffs = [abs(x - s) for x, s in zip(x_size, size)] + diffs = [abs(x - s) for x, s in zip(x_size, size, strict=False)] d1 = [d // 2 for d in diffs] d2 = [d - (d // 2) for d in diffs] @@ -154,71 +158,95 @@ def _pad_crop_img( def pad_img_tensor(x: torch.Tensor, size: Sequence[int]) -> torch.Tensor: - """Pads a tensor + """Pad a tensor of shape (B, C, [Z], Y, X) to the desired spatial dimensions. - Pads a tensor of shape (B, C, [Z], Y, X) to desired spatial dimensions. + Parameters + ---------- + x : torch.Tensor + Input image of shape (B, C, [Z], Y, X). + size : Sequence[int] + Desired spatial size ([Z*], Y*, X*). - Parameters: - ----------- - x (torch.Tensor): Input image of shape (B, C, [Z], Y, X) - size (list or tuple): Desired size ([Z*], Y*, X*) - - Returns: - -------- - The padded tensor + Returns + ------- + torch.Tensor + The padded tensor. """ return _pad_crop_img(x, size, "pad") -def crop_img_tensor(x, size) -> torch.Tensor: - """Crops a tensor. - Crops a tensor of shape (batch, channels, h, w) to a desired height and width - given by a tuple. - Args: - x (torch.Tensor): Input image - size (list or tuple): Desired size (height, width) +def crop_img_tensor(x: torch.Tensor, size: Sequence[int]) -> torch.Tensor: + """Crop a tensor of shape (B, C, [Z], Y, X) to the desired spatial dimensions. + + Parameters + ---------- + x : torch.Tensor + Input image of shape (B, C, [Z], Y, X). + size : Sequence[int] + Desired spatial size ([Z*], Y*, X*). Returns ------- - The cropped tensor + torch.Tensor + The cropped tensor. """ return _pad_crop_img(x, size, "crop") class StableExponential: - """ - Class that redefines the definition of exp() to increase numerical stability. - Naturally, also the definition of log() must change accordingly. - However, it is worth noting that the two operations remain one the inverse of the other, - meaning that x = log(exp(x)) and x = exp(log(x)) are always true. - - Definition: - exp(x) = { - exp(x) if x<=0 - x+1 if x>0 - } - - log(x) = { - x if x<=0 - log(1+x) if x>0 - } + """Numerically stable redefinition of ``exp()`` and its inverse ``log()``. + + The definitions of exp() and log() are redefined to increase numerical stability, + while remaining one the inverse of the other (``x = log(exp(x))`` and + ``x = exp(log(x))`` always hold). + + Definition:: + + exp(x) = { exp(x) if x <= 0 ; x + 1 if x > 0 } + log(x) = { x if x <= 0 ; log(1+x) if x > 0 } NOTE 1: - Within the class everything is done on the tensor given as input to the constructor. - Therefore, when exp() is called, self._tensor.exp() is computed. - When log() is called, torch.log(self._tensor.exp()) is computed instead. + Everything is done on the tensor given as input to the constructor. Therefore, + when exp() is called, ``self._tensor.exp()`` is computed; when log() is called, + ``torch.log(self._tensor.exp())`` is computed instead. NOTE 2: - Given the output from exp(), torch.log() or the log() method of the class give identical results. + Given the output from exp(), ``torch.log()`` or the log() method of the class + give identical results. + + Parameters + ---------- + tensor : torch.Tensor + The tensor on which the stable operations are performed. """ - def __init__(self, tensor): + def __init__(self, tensor: torch.Tensor): + """Constructor. + + Parameters + ---------- + tensor : torch.Tensor + The tensor on which the stable operations are performed. + """ self._raw_tensor = tensor posneg_dic = self.posneg_separation(self._raw_tensor) self.pos_f, self.neg_f = posneg_dic["filter"] self.pos_data, self.neg_data = posneg_dic["value"] - def posneg_separation(self, tensor): + def posneg_separation(self, tensor: torch.Tensor) -> dict: + """Split a tensor into its positive and non-positive parts. + + Parameters + ---------- + tensor : torch.Tensor + The tensor to split. + + Returns + ------- + dict + A dictionary with the positive/negative boolean masks under ``"filter"`` + and the clipped positive/negative tensors under ``"value"``. + """ pos = tensor > 0 pos_tensor = torch.clip(tensor, min=0) @@ -227,32 +255,54 @@ def posneg_separation(self, tensor): return {"filter": [pos, neg], "value": [pos_tensor, neg_tensor]} - def exp(self): + def exp(self) -> torch.Tensor: + """Compute the numerically stable exponential of the tensor. + + Returns + ------- + torch.Tensor + The stable exponential of the input tensor. + """ return torch.exp(self.neg_data) * self.neg_f + (1 + self.pos_data) * self.pos_f - def log(self): + def log(self) -> torch.Tensor: + """Compute the numerically stable logarithm of the tensor. + + Returns + ------- + torch.Tensor + The stable logarithm of the input tensor. + """ return self.neg_data * self.neg_f + torch.log(1 + self.pos_data) * self.pos_f class StableLogVar: - """ - Class that provides a numerically stable implementation of Log-Variance. - Specifically, it uses the exp() and log() formulas defined in `StableExponential` class. + """Numerically stable implementation of Log-Variance. + + It relies on the exp() and log() formulas defined in the `StableExponential` class. + + Parameters + ---------- + logvar : torch.Tensor + The input (true) logvar vector, to be converted in the stable version. + enable_stable : bool, optional + Whether to compute the stable version of log-variance. Default is `True`. + var_eps : float, optional + The minimum value attainable by the variance. Default is `1e-6`. """ def __init__( self, logvar: torch.Tensor, enable_stable: bool = True, var_eps: float = 1e-6 ): - """ - Constructor. + """Constructor. Parameters ---------- - logvar: torch.Tensor - The input (true) logvar vector, to be converted in the Stable version. - enable_stable: bool, optional + logvar : torch.Tensor + The input (true) logvar vector, to be converted in the stable version. + enable_stable : bool, optional Whether to compute the stable version of log-variance. Default is `True`. - var_eps: float, optional + var_eps : float, optional The minimum value attainable by the variance. Default is `1e-6`. """ self._lv = logvar @@ -260,37 +310,59 @@ def __init__( self._eps = var_eps def get(self) -> torch.Tensor: + """Return the (possibly stabilized) log-variance. + + Returns + ------- + torch.Tensor + The log-variance tensor. + """ if self._enable_stable is False: return self._lv return torch.log(self.get_var()) def get_var(self) -> torch.Tensor: - """ - Get Variance from Log-Variance. + """Compute the variance from the log-variance. + + Returns + ------- + torch.Tensor + The variance tensor. """ if self._enable_stable is False: return torch.exp(self._lv) return StableExponential(self._lv).exp() + self._eps def get_std(self) -> torch.Tensor: + """Compute the standard deviation from the log-variance. + + Returns + ------- + torch.Tensor + The standard-deviation tensor. + """ return torch.sqrt(self.get_var()) @property def is_3D(self) -> bool: - """Check if the _lv tensor is 3D. + """Check if the log-variance tensor is 3D. Recall that, in this framework, tensors have shape (B, C, [Z], Y, X). + + Returns + ------- + bool + Whether the tensor is 3D (i.e. has 5 dimensions). """ return self._lv.dim() == 5 def centercrop_to_size(self, size: Sequence[int]) -> None: - """ - Centercrop the log-variance tensor to the desired size. + """Centercrop the log-variance tensor to the desired size. Parameters ---------- - size: torch.Tensor + size : Sequence[int] The desired size of the log-variance tensor. """ assert not self.is_3D, "Centercrop is implemented only for 2D tensors." @@ -304,18 +376,44 @@ def centercrop_to_size(self, size: Sequence[int]) -> None: class StableMean: + """Thin wrapper around a mean tensor exposing stable-distribution helpers. - def __init__(self, mean): + Parameters + ---------- + mean : torch.Tensor + The mean tensor to wrap. + """ + + def __init__(self, mean: torch.Tensor): + """Constructor. + + Parameters + ---------- + mean : torch.Tensor + The mean tensor to wrap. + """ self._mean = mean def get(self) -> torch.Tensor: + """Return the wrapped mean tensor. + + Returns + ------- + torch.Tensor + The mean tensor. + """ return self._mean @property def is_3D(self) -> bool: - """Check if the _mean tensor is 3D. + """Check if the mean tensor is 3D. Recall that, in this framework, tensors have shape (B, C, [Z], Y, X). + + Returns + ------- + bool + Whether the tensor is 3D (i.e. has 5 dimensions). """ return self._mean.dim() == 5 @@ -326,8 +424,8 @@ def centercrop_to_size(self, size: Sequence[int]) -> None: Parameters ---------- - size: torch.Tensor - The desired size of the log-variance tensor. + size : Sequence[int] + The desired size of the mean tensor. """ assert not self.is_3D, "Centercrop is implemented only for 2D tensors." @@ -340,12 +438,37 @@ def centercrop_to_size(self, size: Sequence[int]) -> None: def allow_numpy(func): - """ - All optional arguments are passed as is. positional arguments are checked. if they are numpy array, - they are converted to torch Tensor. + """Wrap a function so that numpy-array positional arguments are cast to tensors. + + Optional (keyword) arguments are passed through unchanged; positional arguments that + are numpy arrays are converted to torch tensors before calling the wrapped function. + + Parameters + ---------- + func : Callable + The function to wrap. + + Returns + ------- + Callable + The wrapped function. """ def numpy_wrapper(*args, **kwargs): + """Cast numpy-array positional arguments to tensors, then call ``func``. + + Parameters + ---------- + *args : Any + Positional arguments; numpy arrays are converted to tensors. + **kwargs : Any + Keyword arguments, passed through unchanged. + + Returns + ------- + Any + The output of the wrapped function. + """ new_args = [] for arg in args: if isinstance(arg, np.ndarray): @@ -360,9 +483,34 @@ def numpy_wrapper(*args, **kwargs): class Interpolate(nn.Module): - """Wrapper for torch.nn.functional.interpolate.""" + """Wrapper for ``torch.nn.functional.interpolate``. + + Parameters + ---------- + size : int or tuple of int, optional + The target output size. Exactly one of `size` and `scale` must be given. + scale : float, optional + The spatial scale factor. Exactly one of `size` and `scale` must be given. + mode : str, optional + The interpolation mode. Default is ``"bilinear"``. + align_corners : bool, optional + The ``align_corners`` flag passed to ``interpolate``. Default is `False`. + """ def __init__(self, size=None, scale=None, mode="bilinear", align_corners=False): + """Constructor. + + Parameters + ---------- + size : int or tuple of int, optional + The target output size. Exactly one of `size` and `scale` must be given. + scale : float, optional + The spatial scale factor. Exactly one of `size` and `scale` must be given. + mode : str, optional + The interpolation mode. Default is ``"bilinear"``. + align_corners : bool, optional + The ``align_corners`` flag passed to ``interpolate``. Default is `False`. + """ super().__init__() assert (size is None) == (scale is not None) self.size = size @@ -370,7 +518,19 @@ def __init__(self, size=None, scale=None, mode="bilinear", align_corners=False): self.mode = mode self.align_corners = align_corners - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Interpolate the input tensor. + + Parameters + ---------- + x : torch.Tensor + The input tensor to interpolate. + + Returns + ------- + torch.Tensor + The interpolated tensor. + """ out = F.interpolate( x, size=self.size, @@ -382,14 +542,24 @@ def forward(self, x): def kl_normal_mc(z, p_mulv, q_mulv): - """ - One-sample estimation of element-wise KL between two diagonal - multivariate normal distributions. Any number of dimensions, + """Estimate the element-wise KL between two diagonal multivariate normals. + + One-sample Monte-Carlo estimation, working for any number of dimensions, with broadcasting supported (be careful). - :param z: - :param p_mulv: - :param q_mulv: - :return: + + Parameters + ---------- + z : torch.Tensor + The sample at which the KL is estimated. + p_mulv : tuple + The (mean, log-variance) wrappers of the prior distribution ``p``. + q_mulv : tuple + The (mean, log-variance) wrappers of the posterior distribution ``q``. + + Returns + ------- + torch.Tensor + The one-sample estimate of the element-wise KL divergence. """ assert isinstance(p_mulv, tuple) assert isinstance(q_mulv, tuple) diff --git a/tests/models/lvae/test_lvae_architecture.py b/tests/models/lvae/test_lvae_architecture.py index 3694e3589..ef6a8487b 100644 --- a/tests/models/lvae/test_lvae_architecture.py +++ b/tests/models/lvae/test_lvae_architecture.py @@ -189,8 +189,6 @@ def test_bottom_up_pass( decoder_conv_strides=decoder_conv_stride, multiscale_count=multiscale_count, ) - first_bottom_up_layer = model.first_bottom_up - lowres_first_bottom_up_layers = model.lowres_first_bottom_ups bottom_up_layers = model.bottom_up_layers assert len(bottom_up_layers) == len( @@ -203,12 +201,7 @@ def test_bottom_up_pass( img_size = model.image_size n_filters = model.n_filters inputs = torch.ones((1, *img_size)) - outputs = model._bottomup_pass( - inp=inputs, - first_bottom_up=first_bottom_up_layer, - lowres_first_bottom_ups=lowres_first_bottom_up_layers, - bottom_up_layers=bottom_up_layers, - ) + outputs = model.bottomup_pass(inputs) exp_img_size = img_size for i in range(len(bottom_up_layers)): if i + 1 > multiscale_count - 1: