diff --git a/CHANGELOG.md b/CHANGELOG.md index cfe16fcf0..8f924dc9c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,9 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm `metric_args={"map": {"backend": ...}}`. - Add multi-scale feature support to the `edgecrafter` model package, exposing the ECViT feature pyramid (strides 8, 16, 32) through the multi-scale feature interface. +- Add multi-channel image input support to the `edgecrafter` (ECViT) model package, like + DINOv2, DINOv3, and TIMM. Pretrained RGB weights are adapted to the requested number + of channels on load. ### Changed diff --git a/docs/source/data/multi_channel.md b/docs/source/data/multi_channel.md index 22e460085..0f620705d 100644 --- a/docs/source/data/multi_channel.md +++ b/docs/source/data/multi_channel.md @@ -40,13 +40,9 @@ The following models support multi-channel image input: | ------------ | ---------------- | ------------------ | | LightlyTrain | DINOv3 | | | LightlyTrain | DINOv2 | | +| LightlyTrain | EdgeCrafter | | | TIMM | All models | [🔗](#models-timm) | -```{note} -EdgeCrafter ECViT models currently support RGB images only and do not support -multi-channel input. -``` - ## Transforms The following image transforms are disabled for images that do not have 3 channels: diff --git a/docs/source/pretrain_distill/models/edgecrafter.md b/docs/source/pretrain_distill/models/edgecrafter.md index b293f8165..02273a397 100644 --- a/docs/source/pretrain_distill/models/edgecrafter.md +++ b/docs/source/pretrain_distill/models/edgecrafter.md @@ -34,11 +34,6 @@ if __name__ == "__main__": See [Distillation method](#methods-distillation) for more details on pretraining and its configuration options. -```{note} -EdgeCrafter ECViT backbones currently support RGB images only. Multi-channel input is -not supported. -``` - ## Supported Models The following ECViT backbone models are supported: diff --git a/src/lightly_train/_models/_model_helpers.py b/src/lightly_train/_models/_model_helpers.py index bd4c620d3..4f12c3a18 100644 --- a/src/lightly_train/_models/_model_helpers.py +++ b/src/lightly_train/_models/_model_helpers.py @@ -18,6 +18,30 @@ logger = logging.getLogger(__name__) +def adjust_conv_input_channels(weight: Tensor, in_chans: int) -> Tensor: + """Adjust a conv weight's input channels to ``in_chans`` (drop or repeat).""" + weights_in_chans = weight.shape[1] + if weights_in_chans > in_chans: + logger.info( + f"Loading pretrained weights with {weights_in_chans} input channels, " + f"but the model has {in_chans} input channels. Keeping only the first " + f"{in_chans} channels of the pretrained weights." + ) + return weight[:, :in_chans, :, :] + if weights_in_chans < in_chans: + logger.info( + f"Loading pretrained weights with {weights_in_chans} input channels, " + f"but the model has {in_chans} input channels. Repeating the channels " + "of the pretrained weights to initialize the extra channels." + ) + repeat_times = in_chans // weights_in_chans + remainder = in_chans % weights_in_chans + weight = weight.repeat(1, repeat_times, 1, 1) + if remainder > 0: + weight = torch.cat([weight, weight[:, :remainder, :, :]], dim=1) + return weight + + def patch_embed_adjust_input_channels_hook( module: Module, state_dict: dict[str, Any], @@ -28,35 +52,29 @@ def patch_embed_adjust_input_channels_hook( """Hook to adjust the number of channels in the state dict to the number of channels in the module. """ - in_chans: Tensor = module.in_chans # type: ignore + in_chans: int = module.in_chans # type: ignore proj_weight_key = f"{prefix}proj.weight" proj_weight = state_dict.get(proj_weight_key) if proj_weight is not None: - weights_in_chans = proj_weight.shape[1] - if weights_in_chans > in_chans: - # Drop last channels - logger.info( - f"Loading pretrained weights with {weights_in_chans} input channels, " - f"but model has {in_chans} input channels. Keeping only the " - f"first {in_chans} channels of the pretrained weights." - ) - proj_weight = proj_weight[:, :in_chans, :, :] - elif weights_in_chans < in_chans: - # Repeat channels to initialize extra channels - logger.info( - f"Loading pretrained weights with {weights_in_chans} input channels, " - f"but model has {in_chans} input channels. Repeating the " - "channels of the pretrained weights to initialize the extra " - "channels." - ) - repeat_times = in_chans // weights_in_chans - remainder = in_chans % weights_in_chans - proj_weight = proj_weight.repeat(1, repeat_times, 1, 1) - if remainder > 0: - proj_weight = torch.cat( - [proj_weight, proj_weight[:, :remainder, :, :]], dim=1 - ) - state_dict[proj_weight_key] = proj_weight + state_dict[proj_weight_key] = adjust_conv_input_channels(proj_weight, in_chans) + + +def conv_pyramid_patch_embed_adjust_input_channels_hook( + module: Module, + state_dict: dict[str, Any], + prefix: str, + *args: Any, + **kwargs: Any, +) -> None: + """Hook to adjust input channels for a conv-pyramid patch embed. + + The input conv is ``convs.0.conv``, not a single ``proj``. + """ + in_chans: int = module.in_chans # type: ignore + conv_weight_key = f"{prefix}convs.0.conv.weight" + conv_weight = state_dict.get(conv_weight_key) + if conv_weight is not None: + state_dict[conv_weight_key] = adjust_conv_input_channels(conv_weight, in_chans) def interpolate_pos_embed_hook( diff --git a/src/lightly_train/_models/ecvit/ecvit.py b/src/lightly_train/_models/ecvit/ecvit.py index 4d9be968c..1ad115933 100644 --- a/src/lightly_train/_models/ecvit/ecvit.py +++ b/src/lightly_train/_models/ecvit/ecvit.py @@ -47,7 +47,9 @@ from torch import Tensor from typing_extensions import Self +from lightly_train import _torch_helpers from lightly_train._activation_checkpointing import maybe_checkpoint +from lightly_train._models import _model_helpers from lightly_train._models.dinov3.dinov3_src.layers.ffn_layers import Mlp from lightly_train._models.dinov3.dinov3_src.layers.rope_position_encoding import ( RopePositionEmbedding, @@ -126,13 +128,20 @@ def apply_rope(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: class ConvPyramidPatchEmbed(nn.Module): - def __init__(self, embed_dim: int = 192, patch_size: int = 16, act: str = "relu"): + def __init__( + self, + embed_dim: int = 192, + patch_size: int = 16, + in_chans: int = 3, + act: str = "relu", + ): super().__init__() if patch_size != 16: raise NotImplementedError( "Only support patch_size=16 for ConvPyramidPatchEmbed." ) + self.in_chans = in_chans num_stages = int(math.log2(patch_size)) - 1 ratios = [2**i for i in range(num_stages, 0, -1)] channels = [embed_dim // r for r in ratios] @@ -142,13 +151,16 @@ def __init__(self, embed_dim: int = 192, patch_size: int = 16, act: str = "relu" ConvNormLayer( # type: ignore[no-untyped-call] in_ch, out_ch, 3, 2, act=act ) - for in_ch, out_ch in zip([3] + channels[:-1], channels) + for in_ch, out_ch in zip([in_chans] + channels[:-1], channels) ] ) self.proj = nn.Conv2d( channels[-1], embed_dim, kernel_size=3, stride=2, padding=1 ) + _torch_helpers.register_load_state_dict_pre_hook( + self, _model_helpers.conv_pyramid_patch_embed_adjust_input_channels_hook + ) def forward(self, x: Tensor) -> Tensor: for conv in self.convs: @@ -172,9 +184,13 @@ def __init__( ) self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) self.num_patches = self.grid_size[0] * self.grid_size[1] + self.in_chans = in_chans self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size ) + _torch_helpers.register_load_state_dict_pre_hook( + self, _model_helpers.patch_embed_adjust_input_channels_hook + ) def forward(self, x: Tensor) -> Tensor: return cast(Tensor, self.proj(x)) @@ -362,7 +378,9 @@ def __init__( embed_dim=embed_dim, ) else: - patch_embed = embed_layer(embed_dim=embed_dim, patch_size=patch_size) + patch_embed = embed_layer( + embed_dim=embed_dim, patch_size=patch_size, in_chans=in_chans + ) self.patch_embed: nn.Module = patch_embed self.patch_size = patch_size @@ -473,6 +491,7 @@ def __init__( embed_dim: int | object = _DEFAULT, num_heads: int | object = _DEFAULT, patch_size: int = 16, + num_input_channels: int = 3, proj_dim: int | None | object = _DEFAULT, num_levels: int = 3, embed_layer: str = "ConvPyramidPatchEmbed", @@ -530,6 +549,7 @@ def __init__( self.name = name self.interaction_indexes = resolved_interaction_indexes self.patch_size = patch_size + self.num_input_channels = num_input_channels self.num_levels = num_levels self.embed_dim = resolved_embed_dim self.num_heads = resolved_num_heads @@ -545,6 +565,7 @@ def __init__( depth=resolved_depth, return_layers=resolved_interaction_indexes, patch_size=patch_size, + in_chans=num_input_channels, embed_layer=EMBED_LAYER_REGISTRY[embed_layer], ffn_layer=FFN_LAYER_REGISTRY[ffn_layer], ffn_ratio=resolved_ffn_ratio, @@ -641,6 +662,13 @@ def multiscale_feature_dims(self) -> list[int]: # All pyramid levels share the same channel dimension. return [self.feature_dim()] * self.num_levels + # TODO(Lionel, 09/26): Levels are not uniformly normalized, so this does not + # fully satisfy the ForwardMultiScaleFeatures contract. When the preset + # `proj_dim` is None (ecvitt, ecvittplus) `forward` projects only the last + # level, leaving levels 0 and 1 as raw interpolated activations. Normalizing + # all levels would change `forward`, which intentionally mirrors upstream + # EdgeCrafter's adapter and is what the LTDETR detection path was tuned + # against. Revisit once a consumer needs comparable statistics across levels. def forward_multiscale_features( self, x: Tensor, layer_indices: Sequence[int] ) -> list[ForwardFeaturesOutput]: diff --git a/src/lightly_train/_models/ecvit/ecvit_package.py b/src/lightly_train/_models/ecvit/ecvit_package.py index d455ac3bb..50b3aed85 100644 --- a/src/lightly_train/_models/ecvit/ecvit_package.py +++ b/src/lightly_train/_models/ecvit/ecvit_package.py @@ -114,15 +114,10 @@ def get_model( ) -> ECViTModelWrapper: """Build an :class:`ECViTModelWrapper` for the given preset. - Multi-channel input is intentionally not supported: ``num_input_channels`` - must be 3. ECViT also does not accept ``model_args`` overrides. + Multi-channel input is supported through ``num_input_channels``: the input + convolution is built for that many channels, and pretrained RGB weights are + adapted to it on load. """ - if num_input_channels != 3: - raise ValueError( - "ECViT backbones only support 3 input channels, got " - f"num_input_channels={num_input_channels}." - ) - model_name = cls.parse_model_name(model_name=model_name) model_info = MODEL_NAME_TO_INFO[model_name] preset_name = model_info["preset_name"] @@ -133,9 +128,11 @@ def get_model( preset_name=preset_name, model_info=model_info ) - return ECViTModelWrapper( - name=preset_name, weights_path=weights_path, **(model_args or {}) - ) + args: dict[str, Any] = {"num_input_channels": num_input_channels} + if model_args is not None: + args.update(model_args) + + return ECViTModelWrapper(name=preset_name, weights_path=weights_path, **args) @classmethod def get_model_wrapper(cls, model: ECViTModelWrapper) -> ECViTModelWrapper: diff --git a/src/lightly_train/_task_models/dinov3_ltdetr/task_model.py b/src/lightly_train/_task_models/dinov3_ltdetr/task_model.py index 72b16282e..42bdbb51f 100644 --- a/src/lightly_train/_task_models/dinov3_ltdetr/task_model.py +++ b/src/lightly_train/_task_models/dinov3_ltdetr/task_model.py @@ -641,14 +641,19 @@ def __init__( self.image_normalize = image_normalize - # Resolve the backbone's expected input channel count. For the DINOv3 - # package we follow the same precedence as DINOV3_PACKAGE.get_model: - # backbone_args["in_chans"] overrides image_normalize, which overrides - # the DINOv3 default of 3. The EdgeCrafter (ECViT) package does not - # support multi-channel input, so we always force 3 there. + # Resolve the backbone's expected input channel count. We follow the same + # precedence as DINOV3_PACKAGE.get_model: backbone_args["in_chans"] + # overrides image_normalize, which overrides the default of 3. package_name = parsed_name["package_name"] + self._expected_input_channels: int + if backbone_args is not None and "in_chans" in backbone_args: + self._expected_input_channels = backbone_args["in_chans"] + elif self.image_normalize is not None: + self._expected_input_channels = len(self.image_normalize["mean"]) + else: + self._expected_input_channels = 3 + if package_name == EDGE_CRAFTER_PACKAGE.name: - self._expected_input_channels: int = 3 # ECViT only supports patch_size=16 (the ECViT-NN uses a # ConvPyramidPatchEmbed that raises NotImplementedError otherwise). # We hard-code it here so the decoder's `config.resolve_auto` @@ -666,12 +671,6 @@ def __init__( "(or set it to 16) to use this model." ) patch_size = 16 - elif backbone_args is not None and "in_chans" in backbone_args: - self._expected_input_channels = backbone_args["in_chans"] - elif self.image_normalize is not None: - self._expected_input_channels = len(self.image_normalize["mean"]) - else: - self._expected_input_channels = 3 # NOTE(Guarin, 08/25): We don't set drop_path_rate=0 here because it is already # set by DINOv3. @@ -696,6 +695,7 @@ def __init__( model_name=parsed_name["backbone_name"], model_args=None, load_weights=load_weights, + **get_model_kwargs, ) ) else: diff --git a/src/lightly_train/_task_models/ltdetr_object_detection/task_model.py b/src/lightly_train/_task_models/ltdetr_object_detection/task_model.py index 594177b5f..39140fbaa 100644 --- a/src/lightly_train/_task_models/ltdetr_object_detection/task_model.py +++ b/src/lightly_train/_task_models/ltdetr_object_detection/task_model.py @@ -34,7 +34,6 @@ ) from lightly_train._models.dinov3.dinov3_vit import DINOv3ViTModelWrapper from lightly_train._models.ecvit.ecvit import ECViTModelWrapper -from lightly_train._models.ecvit.ecvit_package import EDGE_CRAFTER_PACKAGE from lightly_train._pre_post_processing.object_detection import ( ObjectDetectionMetadata, ObjectDetectionOutput, @@ -203,9 +202,7 @@ def __init__( # Resolve the backbone's expected input channel count. # backbone_args["in_chans"] overrides image_normalize, which overrides 3. self._expected_input_channels: int - if package_name == EDGE_CRAFTER_PACKAGE.name: - self._expected_input_channels = 3 - elif backbone_args is not None and "in_chans" in backbone_args: + if backbone_args is not None and "in_chans" in backbone_args: self._expected_input_channels = backbone_args["in_chans"] else: self._expected_input_channels = len(self.image_normalize["mean"]) diff --git a/tests/_models/ecvit/test_ecvit.py b/tests/_models/ecvit/test_ecvit.py index 7d23de684..864016608 100644 --- a/tests/_models/ecvit/test_ecvit.py +++ b/tests/_models/ecvit/test_ecvit.py @@ -12,8 +12,10 @@ import pytest import torch +from torch import nn from lightly_train._models.ecvit import ECVIT_PRESETS, ECViTModelWrapper +from lightly_train._models.ecvit.ecvit import ConvPyramidPatchEmbed from lightly_train._models.model_wrapper import ( ModelWrapper, missing_model_wrapper_attrs, @@ -382,6 +384,112 @@ def test_forward_multiscale_features__matches_forward_features(self) -> None: expected = model.forward_features(image) assert torch.equal(last["features"], expected["features"]) + @pytest.mark.parametrize("num_input_channels", [1, 4, 6]) + def test_forward__multi_channel_input(self, num_input_channels: int) -> None: + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + num_input_channels=num_input_channels, + ) + model.eval() + + with torch.no_grad(): + outputs = model(torch.randn(2, num_input_channels, 32, 32)) + + assert [output.shape for output in outputs] == [ + torch.Size([2, 16, 4, 4]), + torch.Size([2, 16, 2, 2]), + torch.Size([2, 16, 1, 1]), + ] + + def test_init__multi_channel_first_conv(self) -> None: + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + num_input_channels=4, + ) + patch_embed = model.backbone.patch_embed + assert isinstance(patch_embed, ConvPyramidPatchEmbed) + first_conv = patch_embed.convs[0].conv + assert isinstance(first_conv, nn.Conv2d) + assert first_conv.in_channels == 4 + assert first_conv.weight.shape[1] == 4 + + def test_load_weights_path__adapts_input_channels_by_repeat( + self, tmp_path: Path + ) -> None: + # A 3-channel checkpoint loads into a 4-channel model by repeating the first + # convolution's input channels. + source = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + state_dict = source.backbone.state_dict() + weights_path = tmp_path / "ecvitt_rgb.pth" + torch.save(state_dict, weights_path) + + target = ECViTModelWrapper( + name="ecvitt", + weights_path=weights_path, + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + num_input_channels=4, + ) + + source_weight = state_dict["patch_embed.convs.0.conv.weight"] + target_weight = target.backbone.state_dict()["patch_embed.convs.0.conv.weight"] + expected = torch.cat([source_weight, source_weight[:, :1, :, :]], dim=1) + assert target_weight.shape[1] == 4 + assert torch.equal(target_weight, expected) + + def test_load_weights_path__adapts_input_channels_by_slice( + self, tmp_path: Path + ) -> None: + # A 3-channel checkpoint loads into a 1-channel model by keeping the first + # input channel of the first convolution. + source = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + state_dict = source.backbone.state_dict() + weights_path = tmp_path / "ecvitt_rgb.pth" + torch.save(state_dict, weights_path) + + target = ECViTModelWrapper( + name="ecvitt", + weights_path=weights_path, + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + num_input_channels=1, + ) + + source_weight = state_dict["patch_embed.convs.0.conv.weight"] + target_weight = target.backbone.state_dict()["patch_embed.convs.0.conv.weight"] + assert target_weight.shape[1] == 1 + assert torch.equal(target_weight, source_weight[:, :1, :, :]) + @pytest.mark.skipif( os.environ.get("ECVIT_CHECKPOINT_DIR") is None, reason=( diff --git a/tests/_models/test_model_helpers.py b/tests/_models/test_model_helpers.py index b1efa0c76..0f6242af3 100644 --- a/tests/_models/test_model_helpers.py +++ b/tests/_models/test_model_helpers.py @@ -73,3 +73,69 @@ def test__ignores_non_pos_embed_keys(self) -> None: module = _module(torch.zeros(1, 1370, 384)) _model_helpers.interpolate_pos_embed_hook(module, state_dict, "") assert state_dict["blocks.0.attn.qkv.weight"] is weight + + +class _InChansModule(Module): + """Minimal stand-in whose only attribute the input-channel hooks read.""" + + def __init__(self, in_chans: int) -> None: + super().__init__() + self.in_chans = in_chans + + +class TestAdjustConvInputChannels: + def test__noop_same_channels(self) -> None: + weight = torch.randn(8, 3, 3, 3) + out = _model_helpers.adjust_conv_input_channels(weight, 3) + assert torch.equal(out, weight) + + def test__slice_fewer_channels(self) -> None: + weight = torch.randn(8, 3, 3, 3) + out = _model_helpers.adjust_conv_input_channels(weight, 1) + assert out.shape == (8, 1, 3, 3) + assert torch.equal(out, weight[:, :1]) + + def test__repeat_more_channels_with_remainder(self) -> None: + weight = torch.randn(8, 3, 3, 3) + out = _model_helpers.adjust_conv_input_channels(weight, 4) + assert out.shape == (8, 4, 3, 3) + assert torch.equal(out, torch.cat([weight, weight[:, :1]], dim=1)) + + def test__repeat_more_channels_exact_multiple(self) -> None: + weight = torch.randn(8, 3, 3, 3) + out = _model_helpers.adjust_conv_input_channels(weight, 6) + assert out.shape == (8, 6, 3, 3) + assert torch.equal(out, weight.repeat(1, 2, 1, 1)) + + +class TestPatchEmbedAdjustInputChannelsHook: + def test__adjusts_proj_weight(self) -> None: + weight = torch.randn(8, 3, 16, 16) + state_dict = {"patch_embed.proj.weight": weight} + module = _InChansModule(in_chans=5) + _model_helpers.patch_embed_adjust_input_channels_hook( + module, state_dict, "patch_embed." + ) + assert state_dict["patch_embed.proj.weight"].shape == (8, 5, 16, 16) + + def test__ignores_missing_key(self) -> None: + weight = torch.randn(2, 2) + state_dict = {"other.weight": weight} + module = _InChansModule(in_chans=5) + _model_helpers.patch_embed_adjust_input_channels_hook( + module, state_dict, "patch_embed." + ) + assert state_dict["other.weight"] is weight + + +class TestConvPyramidPatchEmbedAdjustInputChannelsHook: + def test__adjusts_first_conv_weight(self) -> None: + weight = torch.randn(8, 3, 3, 3) + state_dict = {"patch_embed.convs.0.conv.weight": weight} + module = _InChansModule(in_chans=2) + _model_helpers.conv_pyramid_patch_embed_adjust_input_channels_hook( + module, state_dict, "patch_embed." + ) + adjusted = state_dict["patch_embed.convs.0.conv.weight"] + assert adjusted.shape == (8, 2, 3, 3) + assert torch.equal(adjusted, weight[:, :2]) diff --git a/tests/_task_models/linear_semantic_segmentation/test_config.py b/tests/_task_models/linear_semantic_segmentation/test_config.py index 6e4aa774d..b5e759436 100644 --- a/tests/_task_models/linear_semantic_segmentation/test_config.py +++ b/tests/_task_models/linear_semantic_segmentation/test_config.py @@ -16,23 +16,26 @@ ) -def test_registry_has_config_for_every_backbone() -> None: - """Every backbone in the DINOv2/DINOv3 packages must have an explicit config. +def test_registry_has_config_for_every_dinov2_backbone() -> None: + """Every DINOv2 backbone must have an explicit config. - Guards against forgetting to register a config when a new backbone is added - to the dinov2 or dinov3 packages. Names that fall through to ``Fallback`` are - derived at runtime and are not considered registered. + DINOv2 configs set ``freeze_mask_token=True`` and ``drop_path_rate=0.0``. + ``Fallback`` keeps ``backbone_name=""``, so ``_check_freeze_mask_token`` does + not fire for it and an unregistered DINOv2 backbone would silently run with + ``freeze_mask_token=False``. Backbones of other packages need no extra fields + and are fine on ``Fallback``. """ missing = [ model_name for model_name in LinearSemanticSegmentation.list_model_names() - if LINEAR_SEG_MODEL_REGISTRY.get( + if model_name.startswith("dinov2/") + and LINEAR_SEG_MODEL_REGISTRY.get( model_name, default=LinearSegConfigRegistry.Fallback ) is LinearSegConfigRegistry.Fallback ] assert not missing, ( - f"Missing explicit LinearSemanticSegmentation configs for backbones: " + f"Missing explicit LinearSemanticSegmentation configs for DINOv2 backbones: " f"{missing}. Add a config for each in LinearSegConfigRegistry." )