From d714812e1bb29983c4c5dd63925d5633a1bea03b Mon Sep 17 00:00:00 2001 From: Szabolcs Botond Lorincz Molnar Date: Sat, 22 Aug 2026 04:40:18 +0200 Subject: [PATCH 1/5] Add multi-scale feature support to edgecrafter ECViT backbone Make ECViTModelWrapper implement MultiScaleFeatureCNN and EdgeCrafterPackage a MultiScaleFeaturePackage, exposing the existing (P3, P4, P5) ECViT pyramid (strides 8, 16, 32) through forward_multiscale_features. Register the four ECViT linear semantic segmentation configs and add unit tests. Part of #928 and #925. --- CHANGELOG.md | 2 + src/lightly_train/_models/ecvit/ecvit.py | 27 ++++- .../_models/ecvit/ecvit_package.py | 4 +- .../linear_semantic_segmentation/config.py | 17 +++ tests/_models/ecvit/test_ecvit.py | 110 ++++++++++++++++++ 5 files changed, 155 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d1119233d..387e6dc41 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,8 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm `pip install "lightly-train[faster-coco-eval]"`; it is then used automatically. Both backends return identical metric values. The backend can be selected explicitly with `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. ### Changed diff --git a/src/lightly_train/_models/ecvit/ecvit.py b/src/lightly_train/_models/ecvit/ecvit.py index 4decb5067..4d9be968c 100644 --- a/src/lightly_train/_models/ecvit/ecvit.py +++ b/src/lightly_train/_models/ecvit/ecvit.py @@ -36,7 +36,7 @@ import math import warnings -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from functools import partial from pathlib import Path from typing import Any, List, cast @@ -57,7 +57,7 @@ ArchitectureInfoGettable, ForwardFeaturesOutput, ForwardPoolOutput, - ModelWrapper, + MultiScaleFeatureCNN, SupportsActivationCheckpointing, ) from lightly_train._task_models.object_detection_components.hybrid_encoder import ( @@ -452,7 +452,7 @@ def forward(self, x: Tensor) -> list[Tensor]: class ECViTModelWrapper( nn.Module, - ModelWrapper, + MultiScaleFeatureCNN, ArchitectureInfoGettable, SupportsActivationCheckpointing, ): @@ -461,6 +461,8 @@ class ECViTModelWrapper( The forward path intentionally follows EdgeCrafter's ECViT adapter: selected ECViT token outputs are averaged, reshaped to a spatial map, interpolated to three levels, projected, and returned as ``(P3, P4, P5)``. + The same pyramid is exposed through the multi-scale feature interface via + ``forward_multiscale_features``. """ def __init__( @@ -631,6 +633,25 @@ def get_model(self) -> Self: def architecture_info(self) -> ArchitectureInfo: return {"model_type": "transformer", "norm_type": "layernorm"} + def multiscale_feature_strides(self) -> list[int]: + # Strides of the three pyramid levels relative to the input (see forward). + return [self.patch_size // 2, self.patch_size, self.patch_size * 2] + + def multiscale_feature_dims(self) -> list[int]: + # All pyramid levels share the same channel dimension. + return [self.feature_dim()] * self.num_levels + + def forward_multiscale_features( + self, x: Tensor, layer_indices: Sequence[int] + ) -> list[ForwardFeaturesOutput]: + if any(index < 0 or index >= self.num_levels for index in layer_indices): + raise ValueError( + f"layer_indices must be in the range [0, {self.num_levels}), but " + f"got {list(layer_indices)}." + ) + features = self.forward(x) + return [{"features": features[index]} for index in layer_indices] + def _load_torch_checkpoint(path: Path) -> object: if not path.exists(): diff --git a/src/lightly_train/_models/ecvit/ecvit_package.py b/src/lightly_train/_models/ecvit/ecvit_package.py index 4f4ec6857..d455ac3bb 100644 --- a/src/lightly_train/_models/ecvit/ecvit_package.py +++ b/src/lightly_train/_models/ecvit/ecvit_package.py @@ -19,7 +19,7 @@ ECViTModelWrapper, ) from lightly_train._models.model_wrapper import ModelWrapper -from lightly_train._models.package import Package +from lightly_train._models.package import MultiScaleFeaturePackage from lightly_train.types import PathLike logger = logging.getLogger(__name__) @@ -72,7 +72,7 @@ class _ECViTModelInfo(TypedDict): } -class EdgeCrafterPackage(Package): +class EdgeCrafterPackage(MultiScaleFeaturePackage): """Package for EdgeCrafter ECViT backbones. The public package name exposed in model strings is ``"edgecrafter"`` diff --git a/src/lightly_train/_task_models/linear_semantic_segmentation/config.py b/src/lightly_train/_task_models/linear_semantic_segmentation/config.py index 03fea4e15..f5e49060a 100644 --- a/src/lightly_train/_task_models/linear_semantic_segmentation/config.py +++ b/src/lightly_train/_task_models/linear_semantic_segmentation/config.py @@ -249,5 +249,22 @@ class DINOv3ConvNextSmallEupe(LinearSemanticSegmentationConfig): class DINOv3ConvNextBaseEupe(LinearSemanticSegmentationConfig): backbone_name: str = "dinov3/convnext-base-eupe" + # --- EdgeCrafter ECViT variants --- + @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvitt-linear") + class EdgeCrafterECViTT(LinearSemanticSegmentationConfig): + backbone_name: str = "edgecrafter/ecvitt" + + @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvittplus-linear") + class EdgeCrafterECViTTPlus(LinearSemanticSegmentationConfig): + backbone_name: str = "edgecrafter/ecvittplus" + + @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvits-linear") + class EdgeCrafterECViTS(LinearSemanticSegmentationConfig): + backbone_name: str = "edgecrafter/ecvits" + + @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvitsplus-linear") + class EdgeCrafterECViTSPlus(LinearSemanticSegmentationConfig): + backbone_name: str = "edgecrafter/ecvitsplus" + class Fallback(LinearSemanticSegmentationConfig): pass # backbone_name stays "", parsed from model_name at runtime diff --git a/tests/_models/ecvit/test_ecvit.py b/tests/_models/ecvit/test_ecvit.py index 9f872d1c9..7d23de684 100644 --- a/tests/_models/ecvit/test_ecvit.py +++ b/tests/_models/ecvit/test_ecvit.py @@ -272,6 +272,116 @@ def test_model_wrapper_interface(self) -> None: assert features["features"].shape == torch.Size([2, 16, 1, 1]) assert pooled["pooled_features"].shape == torch.Size([2, 16, 1, 1]) + def test_multiscale_feature_dims(self) -> None: + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + assert model.multiscale_feature_dims() == [16, 16, 16] + + def test_multiscale_feature_dims__proj_dim_none(self) -> None: + # With proj_dim=None only the last level is projected, but all levels keep + # the same embed_dim channels. + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=None, + ) + assert model.feature_dim() == 16 + assert model.multiscale_feature_dims() == [16, 16, 16] + + def test_multiscale_feature_strides(self) -> None: + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + assert model.multiscale_feature_strides() == [8, 16, 32] + + def test_forward_multiscale_features(self) -> None: + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + model.eval() + with torch.no_grad(): + features = model.forward_multiscale_features( + torch.randn(2, 3, 32, 32), layer_indices=[0, 1, 2] + ) + shapes = [feature["features"].shape for feature in features] + assert shapes == [ + torch.Size([2, 16, 4, 4]), + torch.Size([2, 16, 2, 2]), + torch.Size([2, 16, 1, 1]), + ] + + def test_forward_multiscale_features__order(self) -> None: + # Features are returned in the same order as the requested indices. + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + model.eval() + with torch.no_grad(): + features = model.forward_multiscale_features( + torch.randn(2, 3, 32, 32), layer_indices=[2, 0] + ) + shapes = [feature["features"].shape for feature in features] + assert shapes == [torch.Size([2, 16, 1, 1]), torch.Size([2, 16, 4, 4])] + + @pytest.mark.parametrize("layer_indices", [[-1], [3]]) + def test_forward_multiscale_features__invalid_index( + self, layer_indices: list[int] + ) -> None: + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + with pytest.raises(ValueError, match="layer_indices must be in the range"): + model.forward_multiscale_features( + torch.randn(2, 3, 32, 32), layer_indices=layer_indices + ) + + def test_forward_multiscale_features__matches_forward_features(self) -> None: + # The last level matches forward_features. Eval mode makes the two forward + # passes deterministic. + model = ECViTModelWrapper( + name="ecvitt", + embed_dim=16, + num_heads=1, + depth=1, + interaction_indexes=[0], + proj_dim=16, + ) + model.eval() + image = torch.randn(2, 3, 32, 32) + with torch.no_grad(): + last = model.forward_multiscale_features(image, layer_indices=[2])[0] + expected = model.forward_features(image) + assert torch.equal(last["features"], expected["features"]) + @pytest.mark.skipif( os.environ.get("ECVIT_CHECKPOINT_DIR") is None, reason=( From a11eb8193db64072681164ff3c875dee8bbe1216 Mon Sep 17 00:00:00 2001 From: Szabolcs Botond Lorincz Molnar Date: Sat, 22 Aug 2026 05:36:08 +0200 Subject: [PATCH 2/5] Add multi-channel input support to edgecrafter ECViT backbone Build the ECViT patch embed for num_input_channels and adapt pretrained RGB weights to it on load through a load_state_dict pre-hook, following the DINOv2/DINOv3/TIMM packages. Extend the shared input-channel helper with a conv-pyramid variant, drop the 3-channel restriction in EdgeCrafterPackage, update the docs, and add unit tests. Part of #928. --- CHANGELOG.md | 3 + docs/source/data/multi_channel.md | 6 +- src/lightly_train/_models/_model_helpers.py | 70 +++++++----- src/lightly_train/_models/ecvit/ecvit.py | 27 ++++- .../_models/ecvit/ecvit_package.py | 16 ++- tests/_models/ecvit/test_ecvit.py | 103 ++++++++++++++++++ tests/_models/test_model_helpers.py | 66 +++++++++++ 7 files changed, 248 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 387e6dc41..2139f5e91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,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/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..a56ab9b9e 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, diff --git a/src/lightly_train/_models/ecvit/ecvit_package.py b/src/lightly_train/_models/ecvit/ecvit_package.py index d455ac3bb..43a6339e1 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"] @@ -134,7 +129,10 @@ def get_model( ) return ECViTModelWrapper( - name=preset_name, weights_path=weights_path, **(model_args or {}) + name=preset_name, + weights_path=weights_path, + num_input_channels=num_input_channels, + **(model_args or {}), ) @classmethod diff --git a/tests/_models/ecvit/test_ecvit.py b/tests/_models/ecvit/test_ecvit.py index 7d23de684..e616031e7 100644 --- a/tests/_models/ecvit/test_ecvit.py +++ b/tests/_models/ecvit/test_ecvit.py @@ -382,6 +382,109 @@ 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, + ) + first_conv = model.backbone.patch_embed.convs[0].conv + 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]) From 4bfa230e38d813ba7d50850de959c48901230472 Mon Sep 17 00:00:00 2001 From: Lionel Date: Mon, 21 Sep 2026 16:40:52 +0000 Subject: [PATCH 3/5] chore: apply mdformat to CHANGELOG Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2139f5e91..8f1a7f2df 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,11 +26,11 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm `pip install "lightly-train[faster-coco-eval]"`; it is then used automatically. Both backends return identical metric values. The backend can be selected explicitly with `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. +- 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 From 000f0abcdeeb6350a49b964624a16dd732bdb94d Mon Sep 17 00:00:00 2001 From: Lionel Date: Mon, 21 Sep 2026 17:07:08 +0000 Subject: [PATCH 4/5] fix(ecvit): honor num_input_channels in LTDETR paths, document unnormalized levels Addresses review findings on #950/#951: - ltdetr_object_detection: drop the edgecrafter special case that pinned _expected_input_channels to 3 while get_model was built for len(image_normalize["mean"]). A 4-channel image_normalize built a 4-channel stem but declared 3 in the preprocessor and the ONNX input spec. - dinov3_ltdetr: forward num_input_channels to EDGE_CRAFTER_PACKAGE.get_model and drop the now-false comment claiming ECViT has no multi-channel support. - ecvit_package: merge num_input_channels into model_args instead of passing it alongside **model_args, which raised TypeError when a caller set it in backbone_args. Matches DINOv3Package.get_model. - ecvit: TODO on forward_multiscale_features recording that levels are not uniformly normalized when the preset proj_dim is None. - docs: drop the edgecrafter note contradicting the new multi-channel support. - tests: narrow patch_embed/conv types so mypy passes; reflow CHANGELOG. --- .../pretrain_distill/models/edgecrafter.md | 5 ---- src/lightly_train/_models/ecvit/ecvit.py | 7 ++++++ .../_models/ecvit/ecvit_package.py | 11 ++++----- .../_task_models/dinov3_ltdetr/task_model.py | 24 +++++++++---------- .../ltdetr_object_detection/task_model.py | 5 +--- tests/_models/ecvit/test_ecvit.py | 7 +++++- 6 files changed, 31 insertions(+), 28 deletions(-) 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/ecvit/ecvit.py b/src/lightly_train/_models/ecvit/ecvit.py index a56ab9b9e..1ad115933 100644 --- a/src/lightly_train/_models/ecvit/ecvit.py +++ b/src/lightly_train/_models/ecvit/ecvit.py @@ -662,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 43a6339e1..50b3aed85 100644 --- a/src/lightly_train/_models/ecvit/ecvit_package.py +++ b/src/lightly_train/_models/ecvit/ecvit_package.py @@ -128,12 +128,11 @@ def get_model( preset_name=preset_name, model_info=model_info ) - return ECViTModelWrapper( - name=preset_name, - weights_path=weights_path, - num_input_channels=num_input_channels, - **(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 e616031e7..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, @@ -414,7 +416,10 @@ def test_init__multi_channel_first_conv(self) -> None: proj_dim=16, num_input_channels=4, ) - first_conv = model.backbone.patch_embed.convs[0].conv + 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 From 4baf6a2f0665b7c21829fd75fac72439224f9b6b Mon Sep 17 00:00:00 2001 From: Lionel Date: Mon, 21 Sep 2026 17:12:24 +0000 Subject: [PATCH 5/5] test(linear-seg): narrow registry completeness test to DINOv2 test_registry_has_config_for_every_backbone required an explicit config for every backbone, but 27 of 39 configs only restate backbone_name, which the Fallback path already derives. Only the 12 DINOv2 configs carry real content (freeze_mask_token=True, drop_path_rate=0.0), and Fallback keeps backbone_name="" so _check_freeze_mask_token never fires for it. Narrow the test to DINOv2 backbones and drop the four bare ECViT configs that existed only to satisfy the wider check. --- .../linear_semantic_segmentation/config.py | 17 ----------------- .../linear_semantic_segmentation/test_config.py | 17 ++++++++++------- 2 files changed, 10 insertions(+), 24 deletions(-) diff --git a/src/lightly_train/_task_models/linear_semantic_segmentation/config.py b/src/lightly_train/_task_models/linear_semantic_segmentation/config.py index f5e49060a..03fea4e15 100644 --- a/src/lightly_train/_task_models/linear_semantic_segmentation/config.py +++ b/src/lightly_train/_task_models/linear_semantic_segmentation/config.py @@ -249,22 +249,5 @@ class DINOv3ConvNextSmallEupe(LinearSemanticSegmentationConfig): class DINOv3ConvNextBaseEupe(LinearSemanticSegmentationConfig): backbone_name: str = "dinov3/convnext-base-eupe" - # --- EdgeCrafter ECViT variants --- - @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvitt-linear") - class EdgeCrafterECViTT(LinearSemanticSegmentationConfig): - backbone_name: str = "edgecrafter/ecvitt" - - @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvittplus-linear") - class EdgeCrafterECViTTPlus(LinearSemanticSegmentationConfig): - backbone_name: str = "edgecrafter/ecvittplus" - - @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvits-linear") - class EdgeCrafterECViTS(LinearSemanticSegmentationConfig): - backbone_name: str = "edgecrafter/ecvits" - - @LINEAR_SEG_MODEL_REGISTRY.register("edgecrafter/ecvitsplus-linear") - class EdgeCrafterECViTSPlus(LinearSemanticSegmentationConfig): - backbone_name: str = "edgecrafter/ecvitsplus" - class Fallback(LinearSemanticSegmentationConfig): pass # backbone_name stays "", parsed from model_name at runtime diff --git a/tests/_task_models/linear_semantic_segmentation/test_config.py b/tests/_task_models/linear_semantic_segmentation/test_config.py index 5f29e4942..0e9ce9212 100644 --- a/tests/_task_models/linear_semantic_segmentation/test_config.py +++ b/tests/_task_models/linear_semantic_segmentation/test_config.py @@ -16,22 +16,25 @@ ) -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." )