[WIP] Custom weights support in library by itallix · Pull Request #6982 · open-edge-platform/geti

itallix · 2026-07-06T12:41:56Z

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 backbones and heads into a single weights file

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)

…ification models

…tance_segmentation and keypoints

github-actions · 2026-07-06T12:42:34Z

⏱️ Backend import time — `app.main`

Accelerator	Cumulative (s)	Self (s)
`cpu`	5.423	0.004
`xpu`	4.522	0.003
`cuda`	4.589	0.004

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 · 2026-07-06T12:49:36Z

📊 Test coverage report

Metric	Coverage
Lines	60.2%
Functions	56.9%
Branches	51.8%
Statements	59.8%

github-actions · 2026-07-06T12:52:11Z

🐳 Docker image sizes

Device	Size
`cpu`	4013.1 MB (3.92 GB)
`xpu`	10801.7 MB (10.55 GB)
`cuda`	10208.3 MB (9.97 GB)

kprokofi · 2026-07-07T11:09:38Z

+# Copyright (C) 2026 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+"""Pretrained checkpoint URL registry for detection models.


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

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.

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

kprokofi · 2026-07-07T11:12:53Z

        metric (MetricCallable, optional): Callable for the metric. Defaults to SegmCallable.
        torch_compile (bool, optional): Flag to indicate whether to use torch.compile. Defaults to False.
        tile_config (TileConfig, optional): Configuration for tiling. Defaults to TileConfig(enable_tiler=False).
+        pretrained (bool, optional): Whether to use pretrained model. Defaults to True.


there are no pretrained weights anymore as I can see, we deleted it on the backbone and head level, so we need to reexport full models and store these weights on our S3 storage

Yes, I'm working on it. So far, dinov2-small-seg has been identified as a candidate.

…add urls for segmentation models

codecov-commenter · 2026-07-07T15:00:07Z

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

Codecov Report

❌ Patch coverage is 82.95455% with 45 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../models/classification/utils/pretrained_weights.py	73.80%	22 Missing ⚠️
...itune/backend/lightning/models/detection/deimv2.py	40.00%	6 Missing ⚠️
...ckend/lightning/models/segmentation/dino_v2_seg.py	54.54%	5 Missing ⚠️
...ne/backend/lightning/models/common/rfdetr_mixin.py	87.50%	3 Missing ⚠️
...ckend/lightning/models/utils/pretrained_weights.py	88.88%	3 Missing ⚠️
library/src/getitune/engine/utils/create.py	40.00%	3 Missing ⚠️
library/src/getitune/backend/lightning/engine.py	33.33%	2 Missing ⚠️
...etitune/backend/lightning/models/detection/atss.py	66.66%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

…on to match upgraded weights

itallix added 8 commits July 1, 2026 11:45

feat: load custom weights in engine + enable custom weights for class…

aabdbb0

…ification models

feat: enable custom weights for detection models

3ce1cbe

feat: enable custom weights for instance segmnetation models

aca3039

fix(app): remove cache_dir setup in trainer

5756a0b

fix(vit): use load_checkpoint from utils for custom weights

f30f4f1

fix(library): fix some det model missed pretrained arg

6c9344a

refactor: extract mixin for pretrained weights used in detection, ins…

b1a05ef

…tance_segmentation and keypoints

feat: enable custom weights for segmentation models

e221d60

itallix requested review from a team as code owners July 6, 2026 12:41

Copilot AI review requested due to automatic review settings July 6, 2026 12:41

itallix linked an issue Jul 6, 2026 that may be closed by this pull request

Add option to specify the location of pretrained weights in getitune #5100

Open

itallix marked this pull request as draft July 6, 2026 12:42

github-actions Bot added TEST Any changes in tests Geti Backend Issues related to the Geti application server Geti Library Issues related to Geti Library (OTX) labels Jul 6, 2026

Copilot started reviewing on behalf of itallix July 6, 2026 12:42 View session

Copilot AI reviewed Jul 6, 2026

View reviewed changes

fix: fix pretrained urls and rename module

8b41d7c

kprokofi reviewed Jul 7, 2026

View reviewed changes

itallix force-pushed the vitalii/5100-custom-weights branch from 4b1d26c to 826bee1 Compare July 7, 2026 12:12

test: fix unit tests

659fee2

itallix force-pushed the vitalii/5100-custom-weights branch from 826bee1 to 659fee2 Compare July 7, 2026 12:32

chore: introduce target for pretrained weights (model vs backbone) + …

644cd32

…add urls for segmentation models

itallix force-pushed the vitalii/5100-custom-weights branch from 05c5a7a to 52fcfab Compare July 8, 2026 12:06

refactor: remove pretrained urls for det + inst_seg

3ee72aa

itallix force-pushed the vitalii/5100-custom-weights branch from 52fcfab to 3ee72aa Compare July 8, 2026 12:08

itallix added 2 commits July 8, 2026 16:36

fix: merge weights for dinov2-small-seg & update backbone configurati…

3bb9f5f

…on to match upgraded weights

fix: remap deit keys for vit + fix timm weights loading

bfa12af

Uh oh!

Conversation

itallix commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

How to test

Checklist

Uh oh!

github-actions Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

⏱️ Backend import time — app.main

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📊 Test coverage report

Uh oh!

github-actions Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🐳 Docker image sizes

Uh oh!

kprokofi Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

itallix Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

itallix Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

kprokofi Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

itallix Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

itallix commented Jul 6, 2026 •

edited

Loading

github-actions Bot commented Jul 6, 2026 •

edited

Loading

⏱️ Backend import time — `app.main`

github-actions Bot commented Jul 6, 2026 •

edited

Loading

github-actions Bot commented Jul 6, 2026 •

edited

Loading

codecov-commenter commented Jul 7, 2026 •

edited

Loading