lightly-ai / lightly-ai/lightly-train
Add multi-scale feature support to all backbone packages
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.7k
- Forks
- 116
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 6
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with src/lightly_train/_models/timm/timm.py and the multi-scale protocols in src/lightly_train/_models/model_wrapper.py and src/lightly_train/_models/package.py. Use TIMM's existing forward_intermediates support as the entry point, then add tests for metadata, ordering, shapes, invalid indices, and final-layer consistency. Done means a documented, tested TIMM subset implements the public package and wrapper protocols and rejects unsupported models clearly.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- computer-vision, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 66/100