feat(ecvit): add multi-channel input support

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
 

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:

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:

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(

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]:

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:

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:

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"])