Repository navigation
Add multi-scale feature support to all backbone packages #925
Description
Description
Dense prediction tasks such as semantic segmentation and object detection need normalized intermediate feature maps from multiple backbone layers or stages. Multi-scale features can also be useful during pretraining and distillation, for example to define objectives or transfer representations at more than one level of the network, e.g. wanted in #628 .
LightlyTrain already defines a public interface for this:
ForwardMultiScaleFeatures,MultiScaleFeatureDims, andMultiScaleFeatureStridesMultiScaleFeatureViTandMultiScaleFeatureCNNMultiScaleFeaturePackage
DINOv2 and DINOv3 provide reference implementations:
We would like to extend this interface to the remaining model packages, starting with TIMM. The current TIMMModelWrapper already uses TIMM's forward_intermediates where available, which should provide a useful starting point.
EdgeCrafter is partially supported: ECViTModelWrapper.forward() already returns a three-level feature pyramid, but its wrapper and package do not yet implement the public multi-scale protocols.
Contributions can address one package at a time; there is no need to implement the whole checklist in a single PR. If you would like to work on one, please leave a comment so work is not duplicated.
Expected behavior
For each package:
- Make the package implement
MultiScaleFeaturePackage. - Make its wrapper implement
forward_multiscale_features()and the applicable metadata methods. - Return normalized NCHW feature tensors in the same order as the requested layer or stage indices.
- Expose correct feature dimensions and, depending on the architecture, patch size or feature strides.
- Fail with a clear error for models that cannot expose intermediate features instead of returning incorrect results.
- Add tests for metadata, output ordering and shapes, invalid indices, and consistency with
forward_features()for the final layer or stage.
Package-specific implementations do not have to support every upstream architecture immediately. Clearly documented and tested support for a useful subset is welcome, provided unsupported models are handled explicitly.
Package checklist
- DINOv2
- DINOv3
- TIMM — first implementation target
- EdgeCrafter — adapt the existing feature pyramid to the public interface
- Torchvision
- Ultralytics
- SuperGradients
- RF-DETR
Custom model wrappers are not included in this checklist because they are user-defined, but they can implement the same protocols.
Activity
saud5150 commentedon Aug 3, 2026
Hi @liopeer, I'd like to take TIMM.
I've had a look at the wrapper already —
TIMMModelWrapperpicks betweenforward_intermediates,get_intermediate_layersand plainforward_features, so I'd gate multi-scale support on the same checks. Rough plan:MultiScaleFeatureViT, convolutional ones getMultiScaleFeatureCNN, using the existing_TIMM_ARCH_NAME_PREFIXESlist to tell them apartfeature_info(num_chs/reduction)ValueErrorinstead of wrong outputsOne thing I'd like to check before starting: for ViTs I'd pass
norm=True(same as_forward_intermediatesalready does for the last layer), but CNN stages don't have an equivalent norm, so I'd return those as-is. Does that match what you mean by "normalized"?Aiming for a draft PR within a week. Feel free to assign me.
liopeer commentedon Aug 3, 2026
Yes, that's fine! The rest of the plan also looks correct. I also prefer things failing if the model can't properly implement the interface, instead of relying on ambiguous fallbacks etc.
lorinczszabolcs commentedon Aug 21, 2026
Hi @liopeer, I would like to take the Torchvision one (#927).
I will add multi-scale support to the ResNet and ConvNeXt wrappers:
feature_info.forward_features.Two questions before I open the PR:
conv5at the same stride asstage4(both /32), so it is not a clean extra scale. I plan to leave it unsupported for now and implement ResNet and ConvNeXt. Is that fine, or would you preferstage2/3/4instead?lorinczszabolcs commentedon Aug 22, 2026
Hi @liopeer, I would also like to take the Ultralytics one from the checklist next.
I will add multi-scale support to
UltralyticsModelWrapper:model.model.stride). For each stride I take the last backbone layer that produces it, found with a single dummy forward.forward_features.Classification models and RT-DETR raise a clear error. RT-DETR reports a single stride (
model.stride == [32]) and its decoder exposes no strides, so the stride-based approach cannot find its pyramid reliably. I can revisit RT-DETR separately if you have any suggestions.lorinczszabolcs commentedon Aug 22, 2026
Hi @liopeer,
I took the first item here as well: adapt the existing ECViT feature pyramid to the public multi-scale interface.
What I did:
ECViTModelWrappernow implementsMultiScaleFeatureCNN, andEdgeCrafterPackageis aMultiScaleFeaturePackage.forward_multiscale_featuresexposes the existing(P3, P4, P5)pyramid fromforward. The strides are[8, 16, 32]and all three levels share the same channel dimension.forward_features.One point to confirm:
EdgeCrafterPackagebecoming aMultiScaleFeaturePackagealso makes ECViT a linear semantic segmentation backbone. I added the linear-seg configs for the four ECViT presets, like the other packages. Tell me if you would rather keep ECViT out of linear segmentation for now.I left the second item in #928 (multi-channel input) for a separate change #951.
Add multi-scale feature support to ultralytics YOLO backbones
Add multi-scale feature support to edgecrafter ECViT backbone
lorinczszabolcs commentedon Sep 16, 2026
Hi @liopeer, a quick status on the multi-scale checklist. I have PRs up for four of the packages, all rebased on the latest main:
No rush, just flagging that they are ready for review whenever you have the time.
saud5150 commentedon Sep 20, 2026
Hi @liopeer — following up on the norm question from #939, since I think it's part of something broader that affects everyone on this checklist.
ForwardMultiScaleFeaturesdocuments four guarantees. Onmainthere are three implementations of it and no tests:I wrote a conformance suite against the docstring and ran it on all three. The same four checks fail in each:
ValueErrorValueErrorget_intermediate_layerscollects blocks withif i in blocks_to_take, so results come back in model order rather than requested order. Asking a ConvNeXt for[2, 0, 1]returns shapes[(1,96,56,56), (1,192,28,28), (1,384,14,14)]— stage 0 where stage 2 was requested. Invalid indices surface asAssertionError: only 0 / 1 blocks foundfrom inside the model. Every in-tree caller passes ascending indices, so fixing this changes no current behaviour — but it diverges from what #939, #949, #950 and #972 each implement.On the norm mismatch: I don't think TIMM is the outlier. DINOv3 ConvNeXt builds its per-stage norms as
so stages 0–2 come back raw and only stage 3 is normalized.
Features are normalized.already holds for one level only on the CNN reference path.One more the suite caught:
ConvNeXt(patch_size=...)resizes every stage toh // patch_size, butmultiscale_feature_strides()derives strides from the downsampling convs and reports[4, 8, 16, 32]for stages that are all at stride 16. Thereforward_featuresandforward_multiscale_featuresalso disagree on the last stage —(1,768,7,7)vs(1,768,14,14).Happy to open a PR with the conformance suite (one subclass per package to adopt), the order/duplicate/validation fixes in the three wrappers, and the docstring tightened to match. I'd leave the normalization semantics to you since it's a design call. Does that split work?
feat(ecvit): add multi-scale feature support (#950)
Add multi-scale feature support to ultralytics YOLO backbones