Custom weights support in library

[itallix]

Summary

• Add pretrained flag on model level to control weights loading
• Cleanup pretrained from backbones and heads
• Engine loads weights if pretrained_weights param specified
• Introduce set of mixins for weights loading
• Reviewed app manifests urls
• Merged segmentation backbone and head into a single weights file for DinoV2Seg

How to test

Checklist

• The PR title and description are clear and descriptive
• I have manually tested the changes
• All changes are covered by automated tests
• All related issues are linked to this PR (if applicable)
• Documentation has been updated (if applicable)

[github-actions]

⏱️ Backend import time — app.main

Accelerator	Cumulative (s)	Self (s)
cpu	5.082	0.004
xpu	4.984	0.004
cuda	4.079	0.003

[Copilot]

Pull request overview

This PR refactors pretrained-weight initialization across the library by adding a model-level pretrained flag and routing explicit “base pretrained weights” loading through the Engine (via a new pretrained_weights parameter), while centralizing pretrained URL registries and introducing loader mixins.

Changes:

• Add pretrained flag to Lightning model wrappers and update unit tests to disable implicit weight loading.
• Introduce pretrained URL registries and mixins to standardize how pretrained weights are resolved and loaded.
• Extend engine creation / LightningEngine to optionally load explicit pretrained weights when no full checkpoint is provided.

Reviewed changes

Copilot reviewed 65 out of 65 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
library/tests/unit/backend/lightning/models/instance_segmentation/test_rtmdet_inst.py	Pass pretrained=False in tests to avoid implicit weight loading.
library/tests/unit/backend/lightning/models/instance_segmentation/test_rfdetr_inst.py	Pass pretrained=False in tests; update asserts to use pretrained_urls.
library/tests/unit/backend/lightning/models/instance_segmentation/test_maskrcnn.py	Pass pretrained=False in tests to keep unit tests offline/deterministic.
library/tests/unit/backend/lightning/models/instance_segmentation/test_base.py	Pass pretrained=False; remove create-model test.
library/tests/unit/backend/lightning/models/detection/test_yolox.py	Pass pretrained=False in tests; update call sites accordingly.
library/tests/unit/backend/lightning/models/detection/test_ssd.py	Pass pretrained=False in tests to avoid weight loading.
library/tests/unit/backend/lightning/models/detection/test_rfdetr.py	Pass pretrained=False broadly in RFDETR tests.
library/tests/unit/backend/lightning/models/detection/test_dfine.py	Pass pretrained=False fixture/model construction.
library/tests/unit/backend/lightning/models/detection/test_deimv2.py	Pass pretrained=False; update pretrained URL attribute name.
library/tests/unit/backend/lightning/models/detection/test_deim.py	Pass pretrained=False; update pretrained URL attribute name.
library/tests/unit/backend/lightning/models/classification/backbones/test_timm.py	Remove test that asserted HF cache download behavior.
library/tests/unit/backend/lightning/models/classification/backbones/test_efficientnet.py	Remove explicit pretrained=None parameter usage in tests.
library/src/getitune/engine/utils/create.py	Add pretrained_weights parameter and forward into engine construction.
library/src/getitune/backend/lightning/models/utils/utils.py	Fix typos in TODO comments; shared checkpoint loading utilities remain central.
library/src/getitune/backend/lightning/models/utils/pretrained_weights.py	New mixin to load pretrained weights for models exposing pretrained_urls.
library/src/getitune/backend/lightning/models/keypoint_detection/rtmpose.py	Rename pretrained mapping to pretrained_urls; add pretrained flag; stop loading in _create_model.
library/src/getitune/backend/lightning/models/keypoint_detection/base.py	Add PretrainedWeightsMixin and a pretrained init flag plumbed to base.
library/src/getitune/backend/lightning/models/instance_segmentation/utils/pretrained_urls.py	New centralized pretrained URL registry for instance segmentation.
library/src/getitune/backend/lightning/models/instance_segmentation/rtmdet_inst.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model.
library/src/getitune/backend/lightning/models/instance_segmentation/rfdetr_inst.py	Replace inline URLs with registry; add pretrained flag.
library/src/getitune/backend/lightning/models/instance_segmentation/maskrcnn.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model; remove backbone pretrained config.
library/src/getitune/backend/lightning/models/instance_segmentation/maskrcnn_tv.py	Replace torchvision weights object usage with URL registry and pretrained flag plumbing.
library/src/getitune/backend/lightning/models/instance_segmentation/heads/utils.py	Fix typo in docstring (“backpropagation”).
library/src/getitune/backend/lightning/models/instance_segmentation/base.py	Add PretrainedWeightsMixin and pretrained init flag; remove legacy _create_model that loaded from load_from.
library/src/getitune/backend/lightning/models/detection/yolox.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model.
library/src/getitune/backend/lightning/models/detection/utils/pretrained_urls.py	New centralized pretrained URL registry for detection models.
library/src/getitune/backend/lightning/models/detection/ssd.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model; remove backbone pretrained config.
library/src/getitune/backend/lightning/models/detection/rtdetr.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model.
library/src/getitune/backend/lightning/models/detection/rfdetr.py	Replace inline URLs with registry; add pretrained flag.
library/src/getitune/backend/lightning/models/detection/deimv2.py	Replace inline URLs with registry; add pretrained flag; move key remapping into pretrained_key_mapping.
library/src/getitune/backend/lightning/models/detection/deim.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model.
library/src/getitune/backend/lightning/models/detection/d_fine.py	Replace inline URLs with registry; add pretrained flag; stop loading in _create_model.
library/src/getitune/backend/lightning/models/detection/base.py	Add PretrainedWeightsMixin and pretrained init flag to detection base model.
library/src/getitune/backend/lightning/models/detection/backbones/hgnetv2.py	Remove backbone-level pretrained download logic and embedded URLs/prints.
library/src/getitune/backend/lightning/models/detection/atss.py	Switch to registry; add pretrained flag validation; stop loading in _create_model; remove backbone pretrained config.
library/src/getitune/backend/lightning/models/common/rfdetr_mixin.py	Move pretrained loading into load_pretrained; store internal model/config for later weight alignment.
library/src/getitune/backend/lightning/models/common/backbones/resnet.py	Docstring capitalization fix (“ResNet”).
library/src/getitune/backend/lightning/models/classification/utils/pretrained_urls.py	New centralized pretrained URL registry for classification backbones.
library/src/getitune/backend/lightning/models/classification/utils/load_weights.py	New loader mixins for classification backbones (timm/torchvision/http/pytorchcv/vit).
library/src/getitune/backend/lightning/models/classification/multilabel_models/vit.py	Use shared ViT URLs + loader; add pretrained init flag plumbing; remove inline download logic.
library/src/getitune/backend/lightning/models/classification/multilabel_models/torchvision_model.py	Mix in torchvision loader; add pretrained flag plumbing; docstring tweaks.
library/src/getitune/backend/lightning/models/classification/multilabel_models/timm_model.py	Mix in timm loader; add pretrained flag plumbing; docstring tweaks.
library/src/getitune/backend/lightning/models/classification/multilabel_models/mobilenet_v3.py	Mix in checkpoint loader; add pretrained_urls; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/multilabel_models/efficientnet.py	Mix in pytorchcv loader; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/multilabel_models/base.py	Add pretrained init flag plumbing into the classification base.
library/src/getitune/backend/lightning/models/classification/multiclass_models/vit.py	Use shared ViT URLs + loader; add pretrained init flag plumbing; remove inline download logic.
library/src/getitune/backend/lightning/models/classification/multiclass_models/torchvision_model.py	Mix in torchvision loader; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/multiclass_models/timm_model.py	Mix in timm loader; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/multiclass_models/mobilenet_v3.py	Mix in checkpoint loader; add pretrained_urls; fix variable typo; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/multiclass_models/efficientnet.py	Mix in pytorchcv loader; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/multiclass_models/base.py	Add pretrained init flag plumbing into the classification base.
library/src/getitune/backend/lightning/models/classification/hlabel_models/vit.py	Use shared ViT URLs + loader; add pretrained init flag plumbing; remove inline download logic.
library/src/getitune/backend/lightning/models/classification/hlabel_models/torchvision_model.py	Mix in torchvision loader; add pretrained flag plumbing; docstring tweaks.
library/src/getitune/backend/lightning/models/classification/hlabel_models/timm_model.py	Mix in timm loader; add pretrained flag plumbing; docstring tweaks.
library/src/getitune/backend/lightning/models/classification/hlabel_models/mobilenet_v3.py	Mix in checkpoint loader; add pretrained_urls; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/hlabel_models/efficientnet.py	Mix in pytorchcv loader; add pretrained flag plumbing.
library/src/getitune/backend/lightning/models/classification/hlabel_models/base.py	Add pretrained init flag plumbing into the classification base.
library/src/getitune/backend/lightning/models/classification/backbones/vision_transformer.py	Rename method to load_checkpoint; docstring typo fixes; adjust NPZ loading call.
library/src/getitune/backend/lightning/models/classification/backbones/torchvision.py	Remove backbone-level pretrained handling; store model_name; create unweighted model.
library/src/getitune/backend/lightning/models/classification/backbones/timm.py	Remove backbone-level pretrained flag; create timm model without classifier head.
library/src/getitune/backend/lightning/models/classification/backbones/mobilenet_v3.py	Remove embedded pretrained URL logic; stop downloading in backbone constructor.
library/src/getitune/backend/lightning/models/classification/backbones/efficientnet.py	Remove embedded pretrained URL logic and pytorchcv download calls; docstring typo fix.
library/src/getitune/backend/lightning/models/base.py	Add pretrained flag and base load_pretrained hook called after model creation.
library/src/getitune/backend/lightning/engine.py	Add pretrained_weights param; load pretrained weights when no checkpoint; improve checkpoint-load exception chaining.
application/backend/app/execution/training/getitune_trainer.py	Disable model self-loading (pretrained=False) and pass weights via engine kwargs.

[github-actions]

📊 Test coverage report

Metric	Coverage
Lines	60.5%
Functions	57.0%
Branches	52.3%
Statements	60.1%

[github-actions]

🐳 Docker image sizes

Device	Size
cpu	4013.2 MB (3.92 GB)
xpu	10801.8 MB (10.55 GB)
cuda	10208.4 MB (9.97 GB)

[kprokofi]

Why this file is needed? IMO, It is better to keep this info along with model class for better readability.

[itallix]

Initially, I thought it would be more convenient to handle URL changes across all models in a single place. However, since this likely won't align with our future use cases, I will revert the change.

[itallix]

I've reverted urls for det & inst_seg. But let's keep them separate for classification to avoid duplication between hlabel vs multiclass vs multilabel

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 72.94521% with 79 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../models/classification/utils/pretrained_weights.py	71.05%	22 Missing ⚠️
...ckend/lightning/models/segmentation/dino_v2_seg.py	24.13%	22 Missing ⚠️
...els/classification/backbones/vision_transformer.py	8.69%	21 Missing ⚠️
...itune/backend/lightning/models/detection/deimv2.py	33.33%	6 Missing ⚠️
...kend/lightning/models/common/pretrained_weights.py	88.88%	3 Missing ⚠️
...ne/backend/lightning/models/common/rfdetr_mixin.py	87.50%	3 Missing ⚠️
library/src/getitune/backend/lightning/engine.py	0.00%	1 Missing ⚠️
...etitune/backend/lightning/models/detection/atss.py	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[Copilot]

Pull request overview

Copilot reviewed 89 out of 89 changed files in this pull request and generated 6 comments.

[kprokofi]

My main concern:

1. There is too much incorporated different logic how we load pretrained weights to the models.

I see just two cases: 1) load entire model normally via load_checkpoint 2) load backbone (for classification). We could then define that basic logic (number one) in base model (base.py) and specific logic for backbones to match checkpoint names for classification case in classification base model.

2. We need to define generic logic for storing pretrained URL: either in one file for all models as general registry or near each model definition

3. Logic of the pretrained weights loading better to move to the model level instead of Engine. I know that we agreed it previously differently (as you implemented). However, I've seen here major drawback and logic inconsistency.

[kprokofi]

1. Instead of the Mixin we could just define the "load_pretrained" in the base model. Why do we need to leave empty load_pretrained in base model and define separate default Mixin here for that?

---

Also:

2. we have 2 pretrained_weights.py files under different utils folder that basically cover the same functionality, but for different backend logic. Can we combine them under the same module?

3. utils probably not the best choice for these Mixin like classes that all our models inherited. I propose to place pretrained_weights.py under common/ folder. However, as I wrote above, I would probably avoid defining separate Mixin at all since we have base model class

[itallix]

1. The base class does not match the contract that mixin requests + it should remain agnostic of how weights are fetched, mapped (forcing it there would pollute the base class with assumptions that don't apply to every model)
2. It still makes sense to narrow the scope for CLS specific mixins because classification sits on top of external ecosystems (timm, torchvision, pytorchcv), its loading logic is inherently different + vit has custom parsing logic.
3. Makes sense. Will be moved.

[kprokofi]

IMO, it still looks strange that we define classification models separately in that file and other tasks keep pretrained urls near model classes. I understand that it is to remove duplication between different classification tasks.

Right now the PR moves pretrained URLs into a classification-specific file, while detection/segmentation/etc. still keep URLs next to the model definitions. That creates two different patterns in the codebase, which makes navigation less intuitive.

Would it make sense to keep the interface consistent across tasks? I see two possible options:

1. Keep URLs alongside each task/model (as before). This keeps everything discoverable and allows task-specific overrides when needed.
2. Introduce a common registry (e.g. PRETRAINED_URLS_HUB, or task-specific registries such as PRETRAINED_URLS_CLS_HUB, PRETRAINED_URLS_DETECTION_HUB, etc.), so all pretrained URLs live in one place across the entire project. (that's what you proposed before)

[itallix]

I weighed both of those options, but I think the current setup is a reasonable compromise.

Our main goal here is to fix the duplication across the classification tasks indeed. While this introduces a slight inconsistency with detection and segmentation, it's really just a minor implementation detail of how we initialize class variables, not a major architectural issue.

Staying flexible here lets us solve the immediate problem with duplication. We can always unify the patterns later if we decide to clean it up.

[kprokofi]

please, move these protocol definitions to the top of the module together with other

[itallix]

done

[kprokofi]

I couldn't find that mapping before during loading pretrained weights. Do we really need to rename all weights manually? Why can't we rename layers in VIT backbone.py?

Or simple then upload new weights to match all layers name?

[itallix]

The mapping is there to support loading the official deit format weights specified in our app manifest. It looks like this path was never actually tested before because the library was using .npz + dinov2 format weights for classification and the dinov2 flavor for segmentation. Renaming to adapt for deit means introducing another key_mapping for dinov2.

Not sure if it's a good idea to patch and re-upload the weights, doing the translation at load-time allows us to use official, unmodified upstream deit weights.

@kprokofi I wonder if it would make sense to reevaluate our current ViT implementation in the scope of a separate task. Possibly introduce more specific classes adapted to weight formats that we support?

[kprokofi]

Yes, all VIT can be rewritten using TIMM wrapper directly including DINOV2 I think. Let's consider this as part of TIMM integration task or later after that.

[kprokofi]

hmm, what is then the difference between checkpoint and pretrained_weights?
I could pass same path to checkpoint parameter and load it. It will confuse users.

Also, the logic question arise: if we pass model already being instantiated as LightningModel and pretrained_weights then what does it mean? How we should handle it? In that case as I see it will load the checkpoint via self._model.load_pretrained(pretrained_weights) rewriting the weights of this model.

Probably we need to reconsider scope of responsibility. Probably the model itself should be in charge of managing pretrained_weights instead of Engine.

We could pass pretrained_weights path during model creation in application and introduce this parameter on the model level instead of Engine.

 def __init__(
        self,
        label_info: LabelInfoTypes | int | Sequence,
        data_input_params: DataInputParams | dict | None = None,
        model_name: str = "LightningModel",
        optimizer: OptimizerCallable = DefaultOptimizerCallable,
        scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable,
        metric: MetricCallable = NullMetricCallable,
        torch_compile: bool = False,
        tile_config: TileConfig | dict = TileConfig(enable_tiler=False),
        pretrained: bool = True,
        pretrained_weights: PathLike | None = None
    ) -> None:

I think it make API cleaner

[itallix]

It’s indeed a valid observation and it makes perfect sense to handle it on the model level completely. Probably during our initial discussions, we were biased toward keeping it on the engine level to mirror how checkpoints are handled, but this definitely points us in a different, much cleaner direction. I will refactor that!

[kprokofi]

if pretrained = True and we pass external url to the Engine via new pretrained_weights parameter, the base pretrained weights will be always loaded despite the new url. Please, see my comment in Engine, let's make pretrained_weights is the responsibility area of the model entirely.

[leoll2]

Why depth=2 here?

[itallix]

To hide the stream-wrapper frames and preserve caller provenance, so log locations point to the code producing stdout/stderr rather than LoggerStdoutWriter internals.

[leoll2]

What do you mean that most detectors don't have pretrained weights?

[itallix]

It doesn't make sense. Updated comment.

[leoll2]

What happens if PRETRAINED_WEIGHTS_CACHE_DIR is undefined? The previous code would handle this case:

cache_dir = Path(
    os.environ.get(
        "PRETRAINED_WEIGHTS_CACHE_DIR",
        Path.home() / ".cache" / "torch" / "hub" / "checkpoints",
    )
)

[itallix]

This is always overridden with the default when the library is initialized.

[leoll2]

This comment is misplaced, should be at L26

[itallix]

moved

[leoll2]

Hmm that was not intentional, I may have picked the wrong file. I don't know if the addition of 4 reg tokens makes a big difference in practice, @kprokofi would you advise to switch back to 0 reg?

Btw I verified that both checkpoints (with and without reg) are available in S3.

[kprokofi]

Probably the Ultralytics Interface should be updated as well to fit model level checkpoint loading.
I will consider this change as part of API refactoring task

In general LGTM