diff --git a/src/virtual_stain_flow/models/blocks/blocks.py b/src/virtual_stain_flow/models/blocks/blocks.py index 5071618..457131b 100644 --- a/src/virtual_stain_flow/models/blocks/blocks.py +++ b/src/virtual_stain_flow/models/blocks/blocks.py @@ -31,7 +31,7 @@ Activation layers, commonly used in UNet architectures. """ from abc import ABC, abstractmethod -from typing import Optional +from typing import Optional, Literal import timm import torch.nn as nn @@ -44,8 +44,56 @@ ActivationType ) +PaddingMode = Literal['zeros', 'reflect'] + + +def _validate_padding_mode(padding_mode: PaddingMode) -> None: + """ + Validate the padding mode. + + :param padding_mode: Padding mode to validate. + :raises ValueError: If the padding mode is not 'zeros' or 'reflect'. + """ + if padding_mode not in ('zeros', 'reflect'): + raise ValueError( + f"Unsupported padding_mode: {padding_mode!r}. " + "Expected 'zeros' or 'reflect'." + ) + + +def _is_padded_conv(conv: nn.Conv2d) -> bool: + """ + Check if a convolution layer has padding applied. + + :param conv: Convolution layer to check. + :return: True if the convolution has padding, False otherwise. + """ + if isinstance(conv.padding, str): + return conv.padding != 'valid' + return any(p > 0 for p in conv.padding) + + +def _set_padding_mode_on_padded_convs(module: nn.Module, padding_mode: PaddingMode) -> None: + """ + Set the padding mode on all padded convolution layers within the module. + + :param module: Module containing convolution layers. + :param padding_mode: Padding mode to set ('zeros' or 'reflect'). + """ + if padding_mode == 'zeros': + return + for layer in module.modules(): + if isinstance(layer, nn.Conv2d) and _is_padded_conv(layer): + layer.padding_mode = padding_mode + + class AbstractBlock(ABC, nn.Module): + """ + Abstract base class for all computational blocks. + This class centralizes input type checking and provides common properties + for blocks, such as in_channels, out_channels, and num_units. + """ def __init__( self, in_channels: int, @@ -53,6 +101,15 @@ def __init__( num_units: int = 1, **kwargs: dict ): + """ + Initialize the block with input and output channels, number of units, + and additional keyword arguments. + + :param in_channels: Number of input channels. + :param out_channels: Number of output channels. Defaults to in_channels if None. + :param num_units: Number of units in the block. + :param kwargs: Additional keyword arguments. Not used by the base class. + """ super().__init__() @@ -86,6 +143,9 @@ def __init__( def forward(self, x: Tensor) -> Tensor: """ Forward pass of the block. + + :param x: Input tensor. + :return: Output tensor after applying the block. """ raise NotImplementedError @@ -100,15 +160,33 @@ def out_channels(self) -> int: def num_units(self) -> int: return self._num_units - # These 2 below should be overriden to reflect the actual spatial dimension - # changes the block applies. By default they indicate spatial preserving - # blocks, i.e. the height and width of the input tensor remain unchanged. + """ + Compute the output spatial dimensions of the block. + + By default, the block is assumed to preserve spatial dimensions, + i.e., the output height and width are the same as the input. + + Should be overridden by blocks that change spatial dimensions. + """ def out_h(self, in_h: int) -> int: + """ + Compute the output height given the input height. + + :param in_h: Input height. + :return: Output height after applying the block. + """ return in_h def out_w(self, in_w: int) -> int: + """ + Compute the output width given the input width. + + :param in_w: Input width. + :return: Output width after applying the block. + """ return in_w + class Conv2DConvNeXtBlock(AbstractBlock): """ A ConvNeXt block that applies a sequence of ConvNeXt units @@ -123,7 +201,8 @@ def __init__( num_units: int = 1, norm_type: NormType = 'layer', conv_kernel_size: int = 1, - convnext_kernel_size: int = 7 + convnext_kernel_size: int = 7, + padding_mode: PaddingMode = 'zeros', ): """ Initializes a ConvNeXt block with the specified parameters. @@ -142,7 +221,10 @@ def __init__( :param convnext_kernel_size: Kernel size for the depth-wise convolution layers inside the ConvNeXtBlock. Default is 7, as recommended by Liu et al. (2022), for large receptive fields. + :param padding_mode: Padding mode for padded convolutional layers in + this block. Supported values are 'zeros' and 'reflect'. """ + _validate_padding_mode(padding_mode) out_channels = out_channels or in_channels @@ -165,7 +247,8 @@ def __init__( out_channels=out_channels, kernel_size=conv_kernel_size, stride=1, # fixed - padding='same' # ensure spatial dimensions are preserved + padding='same', # ensure spatial dimensions are preserved + padding_mode=padding_mode, ) ) # timm.models.convnext.ConvNeXtStage adds a normalization @@ -191,8 +274,8 @@ def __init__( # spatial and channel-wise feature extraction computations, # contrasting with the standard Conv2D whose kernel does both # simultaneously. - layers = layers + [ - timm.models.convnext.ConvNeXtBlock( + for _ in range(num_units): + convnext_block = timm.models.convnext.ConvNeXtBlock( in_chs=out_channels, # same input/output channels out_chs=out_channels, kernel_size=convnext_kernel_size, @@ -205,8 +288,9 @@ def __init__( conv_mlp=True, use_grn=True, # GlobalResponseNorm for mlp layers norm_layer=timm.layers.LayerNorm2d, - ) for _ in range(num_units) - ] + ) + _set_padding_mode_on_padded_convs(convnext_block, padding_mode) + layers.append(convnext_block) self.network = nn.Sequential(*layers) @@ -233,7 +317,8 @@ def __init__( out_channels: Optional[int] = None, num_units: int = 1, norm_type: NormType = 'batch', - act_type: ActivationType = 'relu' + act_type: ActivationType = 'relu', + padding_mode: PaddingMode = 'zeros', ): """ Initializes a Conv2D block with the specified parameters. @@ -245,7 +330,10 @@ def __init__( :param norm_type: Type of normalization to apply. Default is 'batch' (GroupNorm with 1 group). :param act_type: Type of activation function to apply. Default is 'relu'. + :param padding_mode: Padding mode for Conv2d layers. Supported values + are 'zeros' and 'reflect'. """ + _validate_padding_mode(padding_mode) out_channels = out_channels or in_channels @@ -271,7 +359,8 @@ def __init__( out_channels=_out, kernel_size=3, # fixed for now stride=1, # fixed - padding='same' # indicate spatial preserving unit + padding='same', # indicate spatial preserving unit + padding_mode=padding_mode, ) ) layers.append( diff --git a/src/virtual_stain_flow/models/unext.py b/src/virtual_stain_flow/models/unext.py index 4fa572f..dc859b2 100644 --- a/src/virtual_stain_flow/models/unext.py +++ b/src/virtual_stain_flow/models/unext.py @@ -25,6 +25,7 @@ PixelShuffle2DUpBlock, ActivationType ) +from .blocks.blocks import PaddingMode, _set_padding_mode_on_padded_convs class ConvNeXtUNet(BaseGeneratorModel): """ @@ -50,6 +51,9 @@ def __init__( pretrained_backbone: bool = False, _num_units: Union[List[int], int] = 2, _pixel_shuffle_preserve_channels: bool = False, + padding_mode: PaddingMode = 'zeros', + _encoder_padding_mode: PaddingMode | None = None, + _decoder_padding_mode: PaddingMode | None = None, ): """ Initializes the ConvNeXtUNet model. @@ -62,20 +66,49 @@ def __init__( :param decoder_compute_block: Type of computation block to use in the decoder. Can be 'convnext' for Conv2DConvNeXtBlock or 'conv2d' for Conv2DNormActBlock. Default is 'convnext'. - :param pretrained_backbone: Whether to use a pretrained ConvNeXtV2_tiny backbone. - Default is False. :param act_type: Type of activation function to use in the output layer. Default is 'sigmoid'. + :param pretrained_backbone: Whether to use a pretrained ConvNeXtV2_tiny backbone. + Default is False. :param _num_units: Number of computation units in each stage. Can be an integer for uniform number of units in all stages, or a list of integers specifying the number of units for each stage. + :param _pixel_shuffle_preserve_channels: Whether to preserve the number of channels + when using PixelShuffle2DUpBlock in the decoder. Default is False. + :param padding_mode: Shared padding mode for padded encoder convolutions + and decoder compute blocks (not upsampling blocks). Supports 'zeros' + (default) and 'reflect'. Reflection + padding requires feature-map dimensions larger than the padding + size (3 for the encoder's 7x7 depthwise convolutions). + :param _encoder_padding_mode: Padding mode used by padded convolution + layers in the ConvNeXtV2 encoder. Supports 'zeros' and 'reflect'. + Defaults to None, inheriting padding_mode; otherwise overrides it. + :param _decoder_padding_mode: Padding mode used by padded convolution + layers in decoder compute blocks. Supports 'zeros' and 'reflect'. + Defaults to None, inheriting padding_mode; otherwise overrides it. """ + if _encoder_padding_mode is None: + _encoder_padding_mode = padding_mode + if _decoder_padding_mode is None: + _decoder_padding_mode = padding_mode + + for mode_name, mode_value in ( + ('padding_mode', padding_mode), + ('_encoder_padding_mode', _encoder_padding_mode), + ('_decoder_padding_mode', _decoder_padding_mode), + ): + if mode_value not in ('zeros', 'reflect'): + raise ValueError( + f"Unsupported {mode_name}: {mode_value!r}. " + "Expected 'zeros' or 'reflect'." + ) super().__init__( in_channels=in_channels, out_channels=out_channels, out_activation=act_type, ) + self._padding_mode = padding_mode # timm implementation of the depth 4 ConvNeXtV2 model # with only down-sampling path, originally optimized @@ -101,6 +134,10 @@ def __init__( self.encoder = convnextv2_model + # Leave unpadded stem, downsampling, and pointwise convolutions unchanged. + _set_padding_mode_on_padded_convs(self.encoder, _encoder_padding_mode) + self._encoder_padding_mode = _encoder_padding_mode + if decoder_up_block == 'pixelshuffle': in_block_handles = [PixelShuffle2DUpBlock] * (depth - 1) in_block_kwargs = [{'preserve_channels': _pixel_shuffle_preserve_channels}] * (depth - 1) @@ -114,6 +151,7 @@ def __init__( ) self._pixel_shuffle_preserve_channels = _pixel_shuffle_preserve_channels self._decoder_up_block = decoder_up_block + self._decoder_padding_mode = _decoder_padding_mode if decoder_compute_block == 'convnext': comp_block_handles = [Conv2DConvNeXtBlock] * (depth - 1) @@ -127,14 +165,25 @@ def __init__( self._decoder_compute_block = decoder_compute_block if isinstance(_num_units, int): - comp_block_kwargs = [{'num_units': _num_units}] * (depth - 1) + comp_block_kwargs = [ + { + 'num_units': _num_units, + 'padding_mode': _decoder_padding_mode, + } + ] * (depth - 1) elif isinstance(_num_units, list): if len(_num_units) != (depth - 1): raise ValueError( f"Expected _num_units to be a list of length {(depth - 1)}, " f"got {len(_num_units)}" ) - comp_block_kwargs = [{'num_units': n} for n in _num_units] + comp_block_kwargs = [ + { + 'num_units': n, + 'padding_mode': _decoder_padding_mode, + } + for n in _num_units + ] else: raise TypeError( f"Expected _num_units to be int or list, " @@ -202,8 +251,11 @@ def to_config(self) -> Dict[str, Any]: "decoder_compute_block": self._decoder_compute_block, "act_type": self._act_type, "pretrained_backbone": self._pretrained_backbone, + "padding_mode": self._padding_mode, "_num_units": self._num_units_cfg, "_pixel_shuffle_preserve_channels": self._pixel_shuffle_preserve_channels, + "_decoder_padding_mode": self._decoder_padding_mode, + "_encoder_padding_mode": self._encoder_padding_mode, }, } @@ -214,12 +266,9 @@ def from_config(cls, config: Dict[str, Any]) -> "ConvNeXtUNet": Accepts either the full dict or just the "init" sub-dict. """ - init_cfg = config.get("init", config) - if "_pixel_shuffle_preserve_channels" not in init_cfg: - # For backward compatibility with configs that don't have this key - init_cfg["_pixel_shuffle_preserve_channels"] = True - if "pretrained_backbone" not in init_cfg: - # For backward compatibility with configs that don't have this key - init_cfg["pretrained_backbone"] = False + init_cfg = dict(config.get("init", config)) + # Preserve historical pixel-shuffle behavior for configs without this key. + init_cfg.setdefault("_pixel_shuffle_preserve_channels", True) + # All other missing options use constructor defaults, including padding. return cls(**init_cfg) diff --git a/tests/conftest.py b/tests/conftest.py index 9c42267..beae89a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -194,9 +194,9 @@ def file_index(tmp_path): for index in range(3): for channel, start_value in (("input_ch1", 1), ("input_ch2", 2), ("target_ch1", 1)): value = start_value + index * (2 if channel.startswith("input") else 1) - Image.fromarray(np.full((10, 10), value, dtype=np.uint16), mode="I;16").save( - paths[channel][index] - ) + Image.fromarray( + np.full((10, 10), value, dtype=np.uint16), + ).save(paths[channel][index]) return pd.DataFrame(paths) diff --git a/tests/datasets/conftest.py b/tests/datasets/conftest.py index 40c5734..ba98096 100644 --- a/tests/datasets/conftest.py +++ b/tests/datasets/conftest.py @@ -51,21 +51,21 @@ def file_index(tmp_path): for idx, path in enumerate(paths["input_ch1"]): value = 1 + idx * 2 # 1, 3, 5 img_array = np.full((10, 10), value, dtype=np.uint16) - img = Image.fromarray(img_array, mode='I;16') + img = Image.fromarray(img_array) img.save(path) # input_ch2: 2, 4, 6 for images 0, 1, 2 for idx, path in enumerate(paths["input_ch2"]): value = 2 + idx * 2 # 2, 4, 6 img_array = np.full((10, 10), value, dtype=np.uint16) - img = Image.fromarray(img_array, mode='I;16') + img = Image.fromarray(img_array) img.save(path) # target_ch1: 1, 2, 3 for images 0, 1, 2 for idx, path in enumerate(paths["target_ch1"]): value = 1 + idx # 1, 2, 3 img_array = np.full((10, 10), value, dtype=np.uint16) - img = Image.fromarray(img_array, mode='I;16') + img = Image.fromarray(img_array) img.save(path) return pd.DataFrame(paths) diff --git a/tests/datasets/ds_engine/test_crop_manifest.py b/tests/datasets/ds_engine/test_crop_manifest.py index b09dee2..66550de 100644 --- a/tests/datasets/ds_engine/test_crop_manifest.py +++ b/tests/datasets/ds_engine/test_crop_manifest.py @@ -23,6 +23,17 @@ ) +def save_test_image(image_array, image_path): + """ + Save an array as an image. + Act as one PIL compatibility boundary. + + :param image_array: Numpy array representing the image. + :param image_path: Path to save the image to. + """ + Image.fromarray(image_array).save(image_path) + + class TestCrop: """Tests for Crop dataclass.""" @@ -251,8 +262,8 @@ def test_update_basic_flow(self, sample_manifest, tmp_path): input_path = tmp_path / 'input.tif' target_path = tmp_path / 'target.tif' - Image.fromarray(input_img, mode='I;16').save(str(input_path)) - Image.fromarray(target_img, mode='I;16').save(str(target_path)) + save_test_image(input_img, input_path) + save_test_image(target_img, target_path) # Create manifest with these files df = pd.DataFrame({ @@ -287,8 +298,8 @@ def test_update_with_offset_crop(self, sample_manifest, tmp_path): input_path = tmp_path / 'input.tif' target_path = tmp_path / 'target.tif' - Image.fromarray(input_img, mode='I;16').save(str(input_path)) - Image.fromarray(target_img, mode='I;16').save(str(target_path)) + save_test_image(input_img, input_path) + save_test_image(target_img, target_path) df = pd.DataFrame({ 'input_ch1': [str(input_path)], @@ -319,8 +330,8 @@ def test_update_multiple_crops_from_same_image(self, sample_manifest, tmp_path): input_path = tmp_path / 'input.tif' target_path = tmp_path / 'target.tif' - Image.fromarray(input_img, mode='I;16').save(str(input_path)) - Image.fromarray(target_img, mode='I;16').save(str(target_path)) + save_test_image(input_img, input_path) + save_test_image(target_img, target_path) df = pd.DataFrame({ 'input_ch1': [str(input_path)], @@ -362,9 +373,9 @@ def test_update_maps_crop_idx_to_manifest_idx(self, sample_manifest, tmp_path): path2 = tmp_path / 'img2.tif' path_target = tmp_path / 'target.tif' - Image.fromarray(img1, mode='I;16').save(str(path1)) - Image.fromarray(img2, mode='I;16').save(str(path2)) - Image.fromarray(np.ones((10, 10), dtype=np.uint16), mode='I;16').save(str(path_target)) + save_test_image(img1, path1) + save_test_image(img2, path2) + save_test_image(np.ones((10, 10), dtype=np.uint16), path_target) # Manifest with two images df = pd.DataFrame({ @@ -400,8 +411,8 @@ def test_update_with_cache(self, sample_manifest, tmp_path): input_path = tmp_path / 'input.tif' target_path = tmp_path / 'target.tif' - Image.fromarray(input_img, mode='I;16').save(str(input_path)) - Image.fromarray(target_img, mode='I;16').save(str(target_path)) + save_test_image(input_img, input_path) + save_test_image(target_img, target_path) df = pd.DataFrame({ 'input_ch1': [str(input_path)], @@ -433,8 +444,8 @@ def test_update_updates_crop_state(self, sample_manifest, tmp_path): input_path = tmp_path / 'input.tif' target_path = tmp_path / 'target.tif' - Image.fromarray(input_img, mode='I;16').save(str(input_path)) - Image.fromarray(target_img, mode='I;16').save(str(target_path)) + save_test_image(input_img, input_path) + save_test_image(target_img, target_path) df = pd.DataFrame({ 'input_ch1': [str(input_path)], @@ -479,8 +490,8 @@ def test_crop_info_property(self, sample_manifest, tmp_path): input_path = tmp_path / 'input.tif' target_path = tmp_path / 'target.tif' - Image.fromarray(input_img, mode='I;16').save(str(input_path)) - Image.fromarray(target_img, mode='I;16').save(str(target_path)) + save_test_image(input_img, input_path) + save_test_image(target_img, target_path) df = pd.DataFrame({ 'input_ch1': [str(input_path)], diff --git a/tests/models/test_compute_blocks.py b/tests/models/test_compute_blocks.py index 8104496..e9cb4f3 100644 --- a/tests/models/test_compute_blocks.py +++ b/tests/models/test_compute_blocks.py @@ -10,6 +10,12 @@ ) +def _has_padding(layer: nn.Conv2d) -> bool: + if isinstance(layer.padding, str): + return layer.padding != "valid" + return any(p > 0 for p in layer.padding) + + class DummyBlock(AbstractBlock): def forward(self, x): return x @@ -89,6 +95,23 @@ def test_convnext_block_exposes_abstract_block_metadata_and_dimensions(self): assert block.out_h(input_tensor.shape[2]) == output.shape[2] assert block.out_w(input_tensor.shape[3]) == output.shape[3] + def test_convnext_block_sets_padding_mode_on_padded_convolutions(self): + block = Conv2DConvNeXtBlock( + in_channels=3, + out_channels=8, + num_units=2, + padding_mode="reflect", + ) + + padded_convolutions = [ + layer + for layer in block.network.modules() + if isinstance(layer, nn.Conv2d) and _has_padding(layer) + ] + + assert padded_convolutions + assert all(layer.padding_mode == "reflect" for layer in padded_convolutions) + class TestConv2DNormActBlock: @pytest.mark.parametrize( @@ -174,3 +197,18 @@ def test_norm_act_block_exposes_abstract_block_metadata_and_dimensions(self): assert block.num_units == 2 assert block.out_h(input_tensor.shape[2]) == output.shape[2] assert block.out_w(input_tensor.shape[3]) == output.shape[3] + + def test_norm_act_block_sets_padding_mode_on_convolutions(self): + block = Conv2DNormActBlock( + in_channels=3, + out_channels=8, + num_units=2, + padding_mode="reflect", + ) + + convolutions = [ + layer for layer in block.network if isinstance(layer, nn.Conv2d) + ] + + assert convolutions + assert all(layer.padding_mode == "reflect" for layer in convolutions) diff --git a/tests/models/test_unext.py b/tests/models/test_unext.py new file mode 100644 index 0000000..ad76544 --- /dev/null +++ b/tests/models/test_unext.py @@ -0,0 +1,112 @@ +from copy import deepcopy + +import pytest +import torch +import torch.nn as nn + +from virtual_stain_flow.models.unext import ConvNeXtUNet + + +def _padded_convolutions(module): + return [ + layer + for layer in module.modules() + if isinstance(layer, nn.Conv2d) + and ( + layer.padding == 'same' + if isinstance(layer.padding, str) + else any(layer.padding) + ) + ] + + +def _assert_padding_modes(model, encoder_mode, decoder_mode): + for branch, expected in ( + (model.encoder, encoder_mode), + *((stage.comp_block, decoder_mode) for stage in model.decoder.stages), + ): + convolutions = _padded_convolutions(branch) + assert convolutions + assert all(layer.padding_mode == expected for layer in convolutions) + + +@pytest.mark.parametrize( + ('options', 'encoder_mode', 'decoder_mode'), + [ + ({}, 'zeros', 'zeros'), + ({'_encoder_padding_mode': 'zeros', '_decoder_padding_mode': 'reflect'}, 'zeros', 'reflect'), + ({'padding_mode': 'reflect'}, 'reflect', 'reflect'), + ({'padding_mode': 'reflect', '_encoder_padding_mode': 'zeros'}, 'zeros', 'reflect'), + ({'padding_mode': 'reflect', '_decoder_padding_mode': 'zeros'}, 'reflect', 'zeros'), + ({'_encoder_padding_mode': 'reflect', '_decoder_padding_mode': 'zeros'}, 'reflect', 'zeros'), + ({'padding_mode': 'reflect', '_encoder_padding_mode': None, '_decoder_padding_mode': None}, 'reflect', 'reflect'), + ], +) +@pytest.mark.parametrize('compute_block', ['convnext', 'conv2d']) +def test_unext_padding_configuration_and_roundtrip( + options, encoder_mode, decoder_mode, compute_block, +): + model = ConvNeXtUNet( + in_channels=1, + out_channels=1, + decoder_compute_block=compute_block, + _num_units=1, + **options, + ) + _assert_padding_modes(model, encoder_mode, decoder_mode) + assert all( + layer.padding_mode == 'zeros' + for layer in model.encoder.modules() + if isinstance(layer, nn.Conv2d) and layer.padding == (0, 0) + ) + + config = model.to_config() + assert config['init']['padding_mode'] == options.get('padding_mode', 'zeros') + assert config['init']['_encoder_padding_mode'] == encoder_mode + assert config['init']['_decoder_padding_mode'] == decoder_mode + restored = ConvNeXtUNet.from_config(config) + _assert_padding_modes(restored, encoder_mode, decoder_mode) + assert restored.to_config() == config + + +@pytest.mark.parametrize('wrapped', [False, True]) +@pytest.mark.parametrize('shared_mode', ['zeros', 'reflect']) +def test_from_config_inherits_shared_padding_without_mutating_input(wrapped, shared_mode): + init = {'in_channels': 1, 'out_channels': 1, '_num_units': 1} + if shared_mode == 'reflect': + init['padding_mode'] = shared_mode + config = {'init': init} if wrapped else init + original = deepcopy(config) + + model = ConvNeXtUNet.from_config(config) + + _assert_padding_modes(model, shared_mode, shared_mode) + assert config == original + + +@pytest.mark.parametrize('name', ['padding_mode', '_encoder_padding_mode', '_decoder_padding_mode']) +@pytest.mark.parametrize('invalid_mode', ['', 'replicate']) +def test_unext_rejects_invalid_padding_modes(name, invalid_mode): + with pytest.raises(ValueError, match=f'Unsupported {name}:'): + ConvNeXtUNet(in_channels=1, out_channels=1, **{name: invalid_mode}) + + +def test_unext_rejects_none_shared_padding_even_with_overrides(): + with pytest.raises(ValueError, match='Unsupported padding_mode:'): + ConvNeXtUNet( + in_channels=1, + out_channels=1, + padding_mode=None, + _encoder_padding_mode='zeros', + _decoder_padding_mode='reflect', + ) + + +def test_unext_shared_reflect_padding_forward(): + model = ConvNeXtUNet( + in_channels=1, out_channels=1, padding_mode='reflect', _num_units=1, + ).eval() + with torch.no_grad(): + output = model(torch.randn(1, 1, 32, 32)) + assert output.shape == (1, 1, 32, 32) + assert torch.isfinite(output).all() diff --git a/tests/vsf_logging/test_mlflow_logger_model_config_unext.py b/tests/vsf_logging/test_mlflow_logger_model_config_unext.py index f9d1cad..e02f611 100644 --- a/tests/vsf_logging/test_mlflow_logger_model_config_unext.py +++ b/tests/vsf_logging/test_mlflow_logger_model_config_unext.py @@ -74,6 +74,9 @@ def test_on_train_start_logs_unext_model_config_and_loss_items( if artifact["content"]["class_path"].endswith("ConvNeXtUNet") ) assert "init" in model_config + assert model_config["init"]["padding_mode"] == "zeros" + assert model_config["init"]["_encoder_padding_mode"] == "zeros" + assert model_config["init"]["_decoder_padding_mode"] == "zeros" assert captured["tags"]["model.0.class_path"].endswith("ConvNeXtUNet")