lightly-ai / lightly-ai/lightly-train

Add multi-scale feature support to all backbone packages

Open
#925 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

distillation enhancement good first issue help wanted models pretraining tracking issue
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:

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.