wehub-resource-sync
16031aae96
CPU tests Workflow / Testing (ubuntu-latest, 3.12) (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.13) (push) Failing after 0s
Mypy Type Check / Type Check (push) Failing after 0s
Docs/Test WorkFlow / Test docs build (push) Failing after 1s
PR Conflict Labeler / labeling (push) Failing after 1s
Dependency resolution / Resolve [tflite] extra — Python 3.12 (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.10) (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.13) (push) Failing after 1s
CPU tests Workflow / build-pkg (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.10) (push) Failing after 0s
CPU tests Workflow / Testing (ubuntu-latest, 3.11) (push) Failing after 0s
Smoke Tests / try-all-models (macos-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (macos-latest, 3.13) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / testing-guardian (push) Has been cancelled
GPU tests Workflow / Testing (push) Has been cancelled
1190 lines
53 KiB
Python
1190 lines
53 KiB
Python
# ------------------------------------------------------------------------
|
|
# RF-DETR
|
|
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
|
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
|
# ------------------------------------------------------------------------
|
|
"""Unit tests for ``rfdetr.models.weights`` — the unified weight-loading and LoRA module.
|
|
|
|
These tests cover ``load_pretrain_weights`` and ``apply_lora`` directly, exercising the unified logic extracted from
|
|
``detr.py`` and ``module_model.py``.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from rfdetr.config import RFDETRBaseConfig, TrainConfig
|
|
from rfdetr.models.weights import _warn_on_partial_load
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_checkpoint(num_classes: int = 91, num_queries: int = 300, group_detr: int = 13) -> dict:
|
|
"""Build a minimal checkpoint dict with the given class count.
|
|
|
|
Args:
|
|
num_classes: Total classes including background (bias shape).
|
|
num_queries: Number of object queries per group.
|
|
group_detr: Number of groups.
|
|
"""
|
|
total_queries = num_queries * group_detr
|
|
state = {
|
|
"class_embed.weight": torch.randn(num_classes, 256),
|
|
"class_embed.bias": torch.randn(num_classes),
|
|
"refpoint_embed.weight": torch.randn(total_queries, 4),
|
|
"query_feat.weight": torch.randn(total_queries, 256),
|
|
"other_layer.weight": torch.randn(10, 10),
|
|
}
|
|
ckpt_args = SimpleNamespace(
|
|
segmentation_head=False,
|
|
patch_size=14,
|
|
class_names=["cat", "dog"],
|
|
)
|
|
return {"model": state, "args": ckpt_args}
|
|
|
|
|
|
def _make_train_config(tmp_path=None) -> TrainConfig:
|
|
"""Return a minimal TrainConfig for use in load_pretrain_weights.
|
|
|
|
Args:
|
|
tmp_path: Optional pytest tmp_path fixture value.
|
|
"""
|
|
return TrainConfig(
|
|
dataset_dir=str(tmp_path / "dataset") if tmp_path else "/nonexistent/dataset",
|
|
output_dir=str(tmp_path / "output") if tmp_path else "/nonexistent/output",
|
|
epochs=10,
|
|
lr=1e-4,
|
|
lr_encoder=1.5e-4,
|
|
batch_size=2,
|
|
weight_decay=1e-4,
|
|
lr_drop=8,
|
|
warmup_epochs=1.0,
|
|
drop_path=0.0,
|
|
multi_scale=False,
|
|
expanded_scales=False,
|
|
do_random_resize_via_padding=False,
|
|
grad_accum_steps=1,
|
|
tensorboard=False,
|
|
)
|
|
|
|
|
|
def _fake_nn_model() -> MagicMock:
|
|
"""Return a MagicMock that behaves enough like an LWDETR nn.Module.
|
|
|
|
Returns:
|
|
MagicMock with reinitialize_detection_head and load_state_dict stubs.
|
|
"""
|
|
model = MagicMock()
|
|
model.reinitialize_detection_head = MagicMock()
|
|
model.load_state_dict = MagicMock()
|
|
return model
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# load_pretrain_weights — reinit scenarios
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLoadPretrainWeightsReinitScenarios:
|
|
"""Verify reinitialize_detection_head call patterns for all class-count scenarios."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_io(self, monkeypatch):
|
|
"""Suppress all download, file-existence, and validation side effects."""
|
|
monkeypatch.setattr("rfdetr.models.weights.download_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_checkpoint_compatibility", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.os.path.isfile", lambda _: True)
|
|
|
|
def test_characterization_fine_tuned_checkpoint_auto_aligns_default_num_classes(self, monkeypatch, tmp_path):
|
|
"""Fine-tuned checkpoint (fewer classes) + default num_classes → 1 reinit to ckpt size.
|
|
|
|
When the user did NOT explicitly set num_classes (default=90), the loader auto-aligns to the checkpoint's class
|
|
count (3 classes = bias shape [3]). Only one reinit fires; no second reinit back to 91.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
|
|
checkpoint = _make_checkpoint(num_classes=3)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
calls = nn_model.reinitialize_detection_head.call_args_list
|
|
assert calls[0] == call(3), f"First reinit must resize to checkpoint size 3, got {calls[0]}"
|
|
assert len(calls) == 1, (
|
|
f"Expected exactly 1 reinit call; got {len(calls)}: {calls}. "
|
|
"A second reinit to 91 would destroy loaded fine-tuned weights."
|
|
)
|
|
assert mc.num_classes == 2, "Auto-aligned checkpoint class count must be persisted back onto ModelConfig."
|
|
|
|
def test_characterization_backbone_pretrain_two_reinits(self, monkeypatch, tmp_path):
|
|
"""Backbone pretrain (more classes in checkpoint) + explicit small num_classes → 2 reinits.
|
|
|
|
Scenario: 91-class COCO checkpoint, user explicitly requested num_classes=2.
|
|
First reinit to 91 so load_state_dict works; second reinit to 3 to match config.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=2)
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
calls = nn_model.reinitialize_detection_head.call_args_list
|
|
assert calls == [call(91), call(3)], f"Expected reinit to [91, 3] (expand then trim), got {calls}"
|
|
|
|
def test_characterization_user_override_larger_than_checkpoint_reexpands(self, monkeypatch, tmp_path):
|
|
"""Explicit num_classes larger than checkpoint → 2 reinits (load then expand back).
|
|
|
|
Scenario: 91-class checkpoint, user explicitly set num_classes=93.
|
|
The head must temporarily align to 91 for loading, then expand back to 94.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=93)
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
calls = nn_model.reinitialize_detection_head.call_args_list
|
|
assert calls == [call(91), call(94)], f"Expected reinit to [91, 94] (load then expand), got {calls}"
|
|
|
|
def test_num_classes_assigned_after_construction_treated_as_explicit(self, monkeypatch, tmp_path):
|
|
"""num_classes assigned post-construction wins over a smaller checkpoint (load then expand).
|
|
|
|
Scenario: config constructed without an explicit num_classes, then ``num_classes`` assigned to 5 — mirroring
|
|
what ``RFDETR._align_num_classes_from_dataset`` does during ``train()`` for a model loaded via
|
|
``from_checkpoint`` (issue #1092). Loading a 3-class checkpoint must align to the checkpoint for loading and
|
|
expand back to 6, NOT auto-align the config back down to the checkpoint's class count.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
|
|
mc.num_classes = 5 # Pydantic records assigned fields in model_fields_set.
|
|
checkpoint = _make_checkpoint(num_classes=3)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
calls = nn_model.reinitialize_detection_head.call_args_list
|
|
assert calls == [call(3), call(6)], f"Expected reinit to [3, 6] (load then expand), got {calls}"
|
|
assert mc.num_classes == 5, "Dataset-aligned num_classes must not be clobbered by the checkpoint."
|
|
|
|
def test_explicit_default_num_classes_treated_as_explicit(self, monkeypatch, tmp_path):
|
|
"""num_classes set explicitly to the class default is honored like any explicit value.
|
|
|
|
Scenario: config constructed with ``num_classes`` equal to the ModelConfig default (so it is
|
|
recorded in ``model_fields_set``), loading a smaller checkpoint. The configured value must be
|
|
preserved — the head aligns to the checkpoint for loading, then expands back to the configured
|
|
size — it must NOT auto-align down to the checkpoint's class count. Guards against
|
|
re-introducing the ``num_classes != default`` clause, which made an explicit default behave
|
|
like "unset" and so silently adopted the checkpoint's class count.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
default_nc = RFDETRBaseConfig.model_fields["num_classes"].default
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=default_nc)
|
|
assert "num_classes" in mc.model_fields_set
|
|
checkpoint = _make_checkpoint(num_classes=3)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
calls = nn_model.reinitialize_detection_head.call_args_list
|
|
assert calls == [call(3), call(default_nc + 1)], (
|
|
f"Expected reinit to [3, {default_nc + 1}] (load at checkpoint, expand back to configured), got {calls}"
|
|
)
|
|
assert mc.num_classes == default_nc, "Explicit default num_classes must be preserved, not auto-aligned."
|
|
|
|
def test_characterization_no_mismatch_no_reinit(self, monkeypatch, tmp_path):
|
|
"""Checkpoint class count matches config → no reinit.
|
|
|
|
Scenario: 91-class checkpoint with num_classes=90. 91 == 90 + 1 → no reinit.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=90)
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
nn_model.reinitialize_detection_head.assert_not_called()
|
|
|
|
def test_keypoint_active_mask_mismatch_is_dropped(self, monkeypatch, tmp_path):
|
|
"""Checkpoint `_kp_active_mask` with mismatched shape is dropped before load_state_dict."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
checkpoint["model"]["_kp_active_mask"] = torch.ones(2, 17, dtype=torch.bool)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
nn_model.state_dict = MagicMock(return_value={"_kp_active_mask": torch.ones(1, 17, dtype=torch.bool)})
|
|
nn_model.load_state_dict.return_value = SimpleNamespace(missing_keys=[], unexpected_keys=[])
|
|
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
loaded_state = nn_model.load_state_dict.call_args[0][0]
|
|
assert "_kp_active_mask" not in loaded_state
|
|
|
|
def test_keypoint_checkpoint_schema_reinitializes_before_and_after_load(
|
|
self,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Schema-dependent keypoint tensors should match checkpoint shape during load, then return to config schema."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
use_grouppose_keypoints=True,
|
|
num_keypoints_per_class=[0, 17],
|
|
)
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
checkpoint["model"]["_kp_active_mask"] = torch.ones(1, 17, dtype=torch.bool)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
nn_model.reinitialize_keypoint_head = MagicMock()
|
|
nn_model.get_num_keypoints_per_class_from_checkpoint = MagicMock(return_value=[17])
|
|
nn_model.state_dict = MagicMock(return_value={"_kp_active_mask": torch.ones(2, 17, dtype=torch.bool)})
|
|
nn_model.load_state_dict.return_value = SimpleNamespace(missing_keys=[], unexpected_keys=[])
|
|
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
assert nn_model.reinitialize_keypoint_head.call_args_list == [call([17]), call([0, 17])]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# load_pretrain_weights — class_names extraction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLoadPretrainWeightsClassNames:
|
|
"""Verify that class_names are extracted from checkpoint and returned."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_io(self, monkeypatch):
|
|
monkeypatch.setattr("rfdetr.models.weights.download_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_checkpoint_compatibility", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.os.path.isfile", lambda _: True)
|
|
|
|
def test_characterization_class_names_extracted_from_checkpoint(self, monkeypatch, tmp_path):
|
|
"""class_names stored in checkpoint args are returned as a list of strings."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=90)
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
checkpoint["args"] = SimpleNamespace(
|
|
segmentation_head=False,
|
|
patch_size=14,
|
|
class_names=["cat", "dog", "bird"],
|
|
)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
result = load_pretrain_weights(nn_model, mc)
|
|
|
|
assert result == ["cat", "dog", "bird"], f"Expected class names from checkpoint, got {result!r}"
|
|
|
|
def test_characterization_empty_class_names_when_absent_from_checkpoint(self, monkeypatch, tmp_path):
|
|
"""Empty list returned when checkpoint has no args or no class_names key."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=90)
|
|
checkpoint = _make_checkpoint(num_classes=91)
|
|
checkpoint.pop("args", None) # no args key at all
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
result = load_pretrain_weights(nn_model, mc)
|
|
|
|
assert result == [], f"Expected empty list when checkpoint has no class_names, got {result!r}"
|
|
|
|
def test_none_pretrain_weights_returns_empty_list_immediately(self, tmp_path):
|
|
"""load_pretrain_weights returns [] without any I/O when pretrain_weights is None."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights=None, device="cpu")
|
|
nn_model = _fake_nn_model()
|
|
|
|
result = load_pretrain_weights(nn_model, mc)
|
|
|
|
assert result == [], f"Expected [] for None pretrain_weights, got {result!r}"
|
|
nn_model.load_state_dict.assert_not_called()
|
|
nn_model.reinitialize_detection_head.assert_not_called()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# load_pretrain_weights — PTL .ckpt format
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLoadPretrainWeightsPTLCkptFormat:
|
|
"""Verify that PTL-native .ckpt checkpoints (state_dict, no model key) are handled."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_io(self, monkeypatch):
|
|
"""Suppress all download, file-existence, and validation side effects."""
|
|
monkeypatch.setattr("rfdetr.models.weights.download_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_checkpoint_compatibility", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.os.path.isfile", lambda _: True)
|
|
|
|
def _make_ptl_checkpoint(
|
|
self,
|
|
num_classes: int = 91,
|
|
num_queries: int = 300,
|
|
group_detr: int = 13,
|
|
) -> dict:
|
|
"""Build a fake PyTorch Lightning (PTL) native checkpoint with state_dict keys prefixed by 'model.'.
|
|
|
|
Args:
|
|
num_classes: Total classes including background (bias shape).
|
|
num_queries: Number of object queries per group.
|
|
group_detr: Number of groups.
|
|
"""
|
|
total_queries = num_queries * group_detr
|
|
raw_state = {
|
|
"class_embed.weight": torch.randn(num_classes, 256),
|
|
"class_embed.bias": torch.randn(num_classes),
|
|
"refpoint_embed.weight": torch.randn(total_queries, 4),
|
|
"query_feat.weight": torch.randn(total_queries, 256),
|
|
"other_layer.weight": torch.randn(10, 10),
|
|
}
|
|
return {
|
|
"state_dict": {f"model.{k}": v for k, v in raw_state.items()},
|
|
"epoch": 10,
|
|
"global_step": 1000,
|
|
}
|
|
|
|
def test_ptl_ckpt_loads_successfully(self, monkeypatch):
|
|
"""PTL .ckpt checkpoints (state_dict without model key) must load without KeyError."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
checkpoint = self._make_ptl_checkpoint(num_classes=91)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
result = load_pretrain_weights(nn_model, mc)
|
|
|
|
nn_model.load_state_dict.assert_called_once()
|
|
assert result == [], f"Expected [] (no args/class_names in checkpoint), got {result!r}"
|
|
|
|
def test_ptl_ckpt_model_prefix_stripped_before_load_state_dict(self, monkeypatch):
|
|
"""Model weights passed to load_state_dict must not carry the 'model.' prefix."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
checkpoint = self._make_ptl_checkpoint(num_classes=91)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
loaded_state = nn_model.load_state_dict.call_args[0][0]
|
|
assert all(not k.startswith("model.") for k in loaded_state), (
|
|
f"Keys passed to load_state_dict must not have 'model.' prefix, got: {list(loaded_state.keys())[:5]}"
|
|
)
|
|
|
|
def test_ptl_ckpt_no_model_prefix_in_state_dict_raises_value_error(self, monkeypatch):
|
|
"""A checkpoint with state_dict but no 'model.'-prefixed keys raises ValueError."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
checkpoint = {"state_dict": {"some_other.key": torch.zeros(1)}, "epoch": 10}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
with pytest.raises(ValueError, match="model\\."):
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
def test_ptl_ckpt_class_names_from_hyper_parameters(self, monkeypatch):
|
|
"""Class names stored in hyper_parameters are returned when args key is absent."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
checkpoint = self._make_ptl_checkpoint(num_classes=91)
|
|
checkpoint["hyper_parameters"] = {"class_names": ["cat", "dog"]}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
result = load_pretrain_weights(nn_model, mc)
|
|
|
|
assert result == ["cat", "dog"], f"Expected class names from hyper_parameters, got {result!r}"
|
|
|
|
def test_ptl_ckpt_args_takes_precedence_over_hyper_parameters(self, monkeypatch):
|
|
"""When both args and hyper_parameters are present, args takes precedence."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
checkpoint = self._make_ptl_checkpoint(num_classes=91)
|
|
checkpoint["args"] = {"class_names": ["from_args"]}
|
|
checkpoint["hyper_parameters"] = {"class_names": ["from_hyper_params"]}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
result = load_pretrain_weights(nn_model, mc)
|
|
|
|
assert result == ["from_args"], f"args must take precedence over hyper_parameters, got {result!r}"
|
|
|
|
def test_ptl_ckpt_non_model_keys_in_state_dict_are_excluded(self, monkeypatch):
|
|
"""Non-model keys in state_dict (optimizer, lr_scheduler) must not appear in checkpoint['model'].
|
|
|
|
Real PTL checkpoints contain keys like 'optimizer.param_groups' and 'lr_scheduler.last_epoch' alongside
|
|
'model.*' weights. The loader must exclude these non-model keys so they do not pollute the state dict passed to
|
|
load_state_dict and do not cause KeyError or unexpected parameter names.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
checkpoint = self._make_ptl_checkpoint(num_classes=91)
|
|
# Inject non-model keys that a real PTL checkpoint would contain
|
|
checkpoint["state_dict"]["optimizer.param_groups"] = torch.zeros(1)
|
|
checkpoint["state_dict"]["lr_scheduler.last_epoch"] = torch.tensor(10)
|
|
checkpoint["state_dict"]["callback_states.ema.shadow_params"] = torch.zeros(4)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
nn_model.load_state_dict.assert_called_once()
|
|
loaded_state = nn_model.load_state_dict.call_args[0][0]
|
|
non_model_keys = [k for k in loaded_state if k.startswith(("optimizer.", "lr_scheduler.", "callback_states."))]
|
|
assert not non_model_keys, f"Non-model keys must be excluded from loaded state; found: {non_model_keys}"
|
|
|
|
def test_ptl_ckpt_torch_compile_orig_mod_prefix_stripped(self, monkeypatch):
|
|
"""PTL .ckpt from a torch.compile-wrapped model must load without KeyError.
|
|
|
|
When a model is wrapped with torch.compile before training, PTL records weights under keys like
|
|
"model._orig_mod.class_embed.bias". The loader must strip both the "model." and the subsequent "_orig_mod."
|
|
segment so the resulting keys match the bare parameter names expected by load_state_dict.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/last.ckpt", device="cpu", num_classes=90)
|
|
raw_state = {
|
|
"class_embed.weight": torch.randn(91, 256),
|
|
"class_embed.bias": torch.randn(91),
|
|
"refpoint_embed.weight": torch.randn(300 * 13, 4),
|
|
"query_feat.weight": torch.randn(300 * 13, 256),
|
|
}
|
|
# Simulate torch.compile: keys are prefixed with "model._orig_mod."
|
|
checkpoint = {
|
|
"state_dict": {f"model._orig_mod.{k}": v for k, v in raw_state.items()},
|
|
"epoch": 5,
|
|
}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
nn_model.load_state_dict.assert_called_once()
|
|
loaded_state = nn_model.load_state_dict.call_args[0][0]
|
|
assert all(not k.startswith(("model.", "_orig_mod.")) for k in loaded_state), (
|
|
f"Keys must have both 'model.' and '_orig_mod.' stripped; got: {list(loaded_state.keys())[:5]}"
|
|
)
|
|
|
|
def test_best_model_callback_format_with_both_model_and_state_dict_still_works(self, monkeypatch):
|
|
"""Checkpoints with both 'model' and 'state_dict' (BestModelCallback format) must still load."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/checkpoint_best_total.pth", device="cpu", num_classes=90)
|
|
# BestModelCallback writes both "model" (raw keys) and "state_dict" (prefixed keys).
|
|
raw_state = {
|
|
"class_embed.weight": torch.randn(91, 256),
|
|
"class_embed.bias": torch.randn(91),
|
|
"refpoint_embed.weight": torch.randn(300 * 13, 4),
|
|
"query_feat.weight": torch.randn(300 * 13, 256),
|
|
}
|
|
checkpoint = {
|
|
"model": raw_state,
|
|
"state_dict": {f"model.{k}": v for k, v in raw_state.items()},
|
|
"epoch": 5,
|
|
}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
nn_model.load_state_dict.assert_called_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# apply_lora
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestApplyLora:
|
|
"""Verify that apply_lora applies LoRA adapters to the backbone encoder.
|
|
|
|
``apply_lora`` lazily imports ``peft`` inside the function body, so we use ``patch.dict("sys.modules", ...)`` to
|
|
intercept the import rather than patching a module-level name.
|
|
"""
|
|
|
|
def test_characterization_apply_lora_wraps_backbone_encoder(self):
|
|
"""apply_lora must call get_peft_model on nn_model.backbone[0].encoder."""
|
|
from rfdetr.models.weights import apply_lora
|
|
|
|
nn_model = MagicMock()
|
|
fake_peft_model = MagicMock()
|
|
|
|
mock_peft = MagicMock()
|
|
mock_peft.get_peft_model.return_value = fake_peft_model
|
|
|
|
with patch.dict("sys.modules", {"peft": mock_peft}):
|
|
apply_lora(nn_model)
|
|
|
|
mock_peft.LoraConfig.assert_called_once()
|
|
lora_kwargs = mock_peft.LoraConfig.call_args.kwargs
|
|
assert lora_kwargs.get("r") == 16, "LoRA rank must be 16"
|
|
assert lora_kwargs.get("lora_alpha") == 16, "LoRA alpha must be 16"
|
|
assert lora_kwargs.get("use_dora") is True, "DoRA must be enabled"
|
|
|
|
assert mock_peft.get_peft_model.call_count == 1, "get_peft_model must be called exactly once"
|
|
assert nn_model.backbone[0].encoder is fake_peft_model, "backbone encoder must be replaced with the peft model"
|
|
|
|
def test_characterization_apply_lora_target_modules(self):
|
|
"""apply_lora must target exactly the 9 expected module names."""
|
|
from rfdetr.models.weights import apply_lora
|
|
|
|
nn_model = MagicMock()
|
|
mock_peft = MagicMock()
|
|
|
|
with patch.dict("sys.modules", {"peft": mock_peft}):
|
|
apply_lora(nn_model)
|
|
|
|
expected_targets = {
|
|
"q_proj",
|
|
"v_proj",
|
|
"k_proj",
|
|
"qkv",
|
|
"query",
|
|
"key",
|
|
"value",
|
|
"cls_token",
|
|
"register_tokens",
|
|
}
|
|
actual_targets = set(mock_peft.LoraConfig.call_args.kwargs.get("target_modules", []))
|
|
assert actual_targets == expected_targets, (
|
|
f"LoRA target_modules mismatch.\nExpected: {expected_targets}\nGot: {actual_targets}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-group query embedding slicing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _labelled_query_tensor(num_queries: int, group_detr: int, dim: int = 2) -> torch.Tensor:
|
|
"""Build a query embedding tensor where row ``g * num_queries + q`` encodes ``[g * 100 + q, 0, ...]``.
|
|
|
|
This lets tests check the per-group ordering of the result without floating-point
|
|
fuzz: the first column carries the (group, query) identity directly.
|
|
"""
|
|
rows = []
|
|
for g in range(group_detr):
|
|
for q in range(num_queries):
|
|
rows.append([float(g * 100 + q)] + [0.0] * (dim - 1))
|
|
return torch.tensor(rows, dtype=torch.float32)
|
|
|
|
|
|
class TestSliceQueryParamPerGroup:
|
|
"""Direct unit tests for ``_slice_query_param_per_group``.
|
|
|
|
The helper is the fix for a latent bug where a flat ``tensor[:N]`` slice scrambled per-group structure when
|
|
``num_queries`` decreased with ``group_detr > 1``. See the docstring in ``rfdetr.models.weights`` for the
|
|
``LWDETR`` packing layout that motivates these tests.
|
|
"""
|
|
|
|
def test_returns_input_unchanged_when_dimensions_match(self):
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = _labelled_query_tensor(num_queries=4, group_detr=3)
|
|
out = _slice_query_param_per_group(tensor, 4, 3, target_num_queries=4, target_group_detr=3)
|
|
assert out is tensor
|
|
|
|
def test_num_queries_decrease_preserves_per_group_structure(self):
|
|
"""The bug being fixed: 4→2 queries with 3 groups must keep first 2 of each group."""
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = _labelled_query_tensor(num_queries=4, group_detr=3)
|
|
out = _slice_query_param_per_group(tensor, 4, 3, target_num_queries=2, target_group_detr=3)
|
|
# Expect rows: g0q0, g0q1, g1q0, g1q1, g2q0, g2q1 → labels 0, 1, 100, 101, 200, 201.
|
|
labels = out[:, 0].int().tolist()
|
|
assert labels == [0, 1, 100, 101, 200, 201], (
|
|
f"Per-group structure scrambled. A flat slice would give {tensor[:6, 0].int().tolist()}."
|
|
)
|
|
|
|
def test_group_detr_decrease_drops_tail_groups(self):
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = _labelled_query_tensor(num_queries=4, group_detr=3)
|
|
out = _slice_query_param_per_group(tensor, 4, 3, target_num_queries=4, target_group_detr=2)
|
|
labels = out[:, 0].int().tolist()
|
|
# First 2 groups, all 4 queries each.
|
|
assert labels == [0, 1, 2, 3, 100, 101, 102, 103]
|
|
|
|
def test_both_decrease(self):
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = _labelled_query_tensor(num_queries=4, group_detr=3)
|
|
out = _slice_query_param_per_group(tensor, 4, 3, target_num_queries=2, target_group_detr=2)
|
|
labels = out[:, 0].int().tolist()
|
|
assert labels == [0, 1, 100, 101]
|
|
|
|
def test_falls_back_to_flat_slice_on_inconsistent_shape(self):
|
|
"""If args don't match the tensor's flat length, defer to legacy behavior."""
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
weird = torch.arange(7, dtype=torch.float32).unsqueeze(1) # shape [7, 1], not 4*3=12
|
|
out = _slice_query_param_per_group(weird, 4, 3, target_num_queries=2, target_group_detr=2)
|
|
# Legacy: tensor[:4]
|
|
assert out.shape == (4, 1)
|
|
assert out[:, 0].tolist() == [0.0, 1.0, 2.0, 3.0]
|
|
|
|
@pytest.mark.parametrize(
|
|
"ckpt_nq,ckpt_g,tgt_nq,tgt_g,expected_labels",
|
|
[
|
|
pytest.param(
|
|
4,
|
|
3,
|
|
8,
|
|
3,
|
|
[0, 1, 2, 3, 100, 101, 102, 103, 200, 201, 202, 203],
|
|
id="nq_expands_g_equal",
|
|
),
|
|
pytest.param(
|
|
4,
|
|
2,
|
|
4,
|
|
4,
|
|
[0, 1, 2, 3, 100, 101, 102, 103],
|
|
id="g_expands_nq_equal",
|
|
),
|
|
pytest.param(
|
|
4,
|
|
3,
|
|
8,
|
|
2,
|
|
[0, 1, 2, 3, 100, 101, 102, 103],
|
|
id="nq_expands_g_shrinks",
|
|
),
|
|
pytest.param(
|
|
4,
|
|
3,
|
|
2,
|
|
4,
|
|
[0, 1, 100, 101, 200, 201],
|
|
id="nq_shrinks_g_expands",
|
|
),
|
|
pytest.param(
|
|
4,
|
|
3,
|
|
8,
|
|
4,
|
|
[0, 1, 2, 3, 100, 101, 102, 103, 200, 201, 202, 203],
|
|
id="both_expand",
|
|
),
|
|
],
|
|
)
|
|
def test_expansion_combos(
|
|
self,
|
|
ckpt_nq: int,
|
|
ckpt_g: int,
|
|
tgt_nq: int,
|
|
tgt_g: int,
|
|
expected_labels: list[int],
|
|
) -> None:
|
|
"""Min(target, ckpt) along each axis produces the correct per-group prefix."""
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = _labelled_query_tensor(num_queries=ckpt_nq, group_detr=ckpt_g)
|
|
out = _slice_query_param_per_group(tensor, ckpt_nq, ckpt_g, tgt_nq, tgt_g)
|
|
assert out[:, 0].int().tolist() == expected_labels
|
|
|
|
def test_num_queries_expansion_returns_smaller_tensor(self):
|
|
"""When target > ckpt, return min-per-group; load_state_dict will reject."""
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = _labelled_query_tensor(num_queries=4, group_detr=3)
|
|
out = _slice_query_param_per_group(tensor, 4, 3, target_num_queries=8, target_group_detr=3)
|
|
# min(4, 8) = 4 per group, all 3 groups → 12 rows == input length.
|
|
assert out.shape == (12, 2)
|
|
|
|
@pytest.mark.parametrize(
|
|
"ckpt_nq,ckpt_g,tgt_nq,tgt_g",
|
|
[
|
|
pytest.param(0, 3, 2, 3, id="ckpt_nq=0"),
|
|
pytest.param(-1, 3, 2, 3, id="ckpt_nq=-1"),
|
|
pytest.param(4, 0, 2, 3, id="ckpt_g=0"),
|
|
pytest.param(4, -1, 2, 3, id="ckpt_g=-1"),
|
|
pytest.param(4, 3, 0, 3, id="tgt_nq=0"),
|
|
pytest.param(4, 3, -1, 3, id="tgt_nq=-1"),
|
|
pytest.param(4, 3, 2, 0, id="tgt_g=0"),
|
|
pytest.param(4, 3, 2, -1, id="tgt_g=-1"),
|
|
],
|
|
)
|
|
def test_raises_on_non_positive_dimension(self, ckpt_nq: int, ckpt_g: int, tgt_nq: int, tgt_g: int) -> None:
|
|
"""ValueError raised when any dimension arg is zero or negative."""
|
|
from rfdetr.models.weights import _slice_query_param_per_group
|
|
|
|
tensor = torch.zeros(12, 2)
|
|
with pytest.raises(ValueError, match="must be positive"):
|
|
_slice_query_param_per_group(tensor, ckpt_nq, ckpt_g, tgt_nq, tgt_g)
|
|
|
|
|
|
class TestLoadPretrainWeightsPerGroupQuerySlice:
|
|
"""End-to-end check that ``load_pretrain_weights`` invokes per-group slicing."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_io(self, monkeypatch):
|
|
monkeypatch.setattr("rfdetr.models.weights.download_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_pretrain_weights", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.validate_checkpoint_compatibility", lambda *a, **kw: None)
|
|
monkeypatch.setattr("rfdetr.models.weights.os.path.isfile", lambda _: True)
|
|
|
|
def _make_args_dict_checkpoint(self, num_queries: int, group_detr: int) -> dict:
|
|
"""Build a checkpoint with labelled query weights and dict-style args."""
|
|
labelled_refpoint = _labelled_query_tensor(num_queries, group_detr, dim=4)
|
|
labelled_query_feat = _labelled_query_tensor(num_queries, group_detr, dim=256)
|
|
state = {
|
|
"class_embed.weight": torch.randn(91, 256),
|
|
"class_embed.bias": torch.randn(91),
|
|
"refpoint_embed.weight": labelled_refpoint,
|
|
"query_feat.weight": labelled_query_feat,
|
|
}
|
|
# Dict-style args payload used to exercise the checkpoint-loading path.
|
|
return {"model": state, "args": {"num_queries": num_queries, "group_detr": group_detr}}
|
|
|
|
def test_decreasing_num_queries_preserves_per_group_structure(self, monkeypatch, tmp_path):
|
|
"""Real flow: checkpoint(nq=4, g=3) → model(nq=2, g=3).
|
|
|
|
Group structure must be preserved.
|
|
"""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
num_queries=2,
|
|
num_select=2,
|
|
group_detr=3,
|
|
)
|
|
checkpoint = self._make_args_dict_checkpoint(num_queries=4, group_detr=3)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
passed_state = nn_model.load_state_dict.call_args[0][0]
|
|
refpoint = passed_state["refpoint_embed.weight"]
|
|
query_feat = passed_state["query_feat.weight"]
|
|
expected = [0, 1, 100, 101, 200, 201]
|
|
# First column carries (group, query) identity (see _labelled_query_tensor).
|
|
assert refpoint[:, 0].int().tolist() == expected, (
|
|
"Per-group structure was not preserved in refpoint_embed.weight."
|
|
)
|
|
assert query_feat[:, 0].int().tolist() == expected, (
|
|
"Per-group structure was not preserved in query_feat.weight."
|
|
)
|
|
|
|
def test_legacy_checkpoint_without_args_falls_back_to_flat_slice(self, monkeypatch, tmp_path):
|
|
"""No ``args`` in checkpoint → preserve the legacy flat slice (backward compat)."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
num_queries=2,
|
|
num_select=2,
|
|
group_detr=3,
|
|
)
|
|
checkpoint = self._make_args_dict_checkpoint(num_queries=4, group_detr=3)
|
|
del checkpoint["args"] # legacy
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
passed_state = nn_model.load_state_dict.call_args[0][0]
|
|
refpoint = passed_state["refpoint_embed.weight"]
|
|
query_feat = passed_state["query_feat.weight"]
|
|
# Legacy flat slice: first 2*3=6 rows of the original 12. Original rows
|
|
# are labelled 0,1,2,3,100,101,102,103,200,201,202,203 → first 6 are
|
|
# 0,1,2,3,100,101.
|
|
expected = [0, 1, 2, 3, 100, 101]
|
|
assert refpoint[:, 0].int().tolist() == expected
|
|
assert query_feat[:, 0].int().tolist() == expected
|
|
|
|
def test_decreasing_group_detr_drops_tail_groups(self, monkeypatch, tmp_path):
|
|
"""Checkpoint(nq=4, g=3) → model(nq=4, g=2): tail group dropped, retained groups intact."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
num_queries=4,
|
|
num_select=4,
|
|
group_detr=2,
|
|
)
|
|
checkpoint = self._make_args_dict_checkpoint(num_queries=4, group_detr=3)
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
passed_state = nn_model.load_state_dict.call_args[0][0]
|
|
refpoint = passed_state["refpoint_embed.weight"]
|
|
query_feat = passed_state["query_feat.weight"]
|
|
expected = [0, 1, 2, 3, 100, 101, 102, 103]
|
|
assert refpoint[:, 0].int().tolist() == expected
|
|
assert query_feat[:, 0].int().tolist() == expected
|
|
|
|
def test_decreasing_num_queries_namespace_args(self, monkeypatch, tmp_path):
|
|
"""Namespace-style args in checkpoint trigger per-group slice identical to dict-style."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
num_queries=2,
|
|
num_select=2,
|
|
group_detr=3,
|
|
)
|
|
labelled_refpoint = _labelled_query_tensor(num_queries=4, group_detr=3, dim=4)
|
|
labelled_query_feat = _labelled_query_tensor(num_queries=4, group_detr=3, dim=256)
|
|
checkpoint = {
|
|
"model": {
|
|
"class_embed.weight": torch.randn(91, 256),
|
|
"class_embed.bias": torch.randn(91),
|
|
"refpoint_embed.weight": labelled_refpoint,
|
|
"query_feat.weight": labelled_query_feat,
|
|
},
|
|
"args": SimpleNamespace(num_queries=4, group_detr=3),
|
|
}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
passed_state = nn_model.load_state_dict.call_args[0][0]
|
|
refpoint = passed_state["refpoint_embed.weight"]
|
|
query_feat = passed_state["query_feat.weight"]
|
|
expected = [0, 1, 100, 101, 200, 201]
|
|
assert refpoint[:, 0].int().tolist() == expected
|
|
assert query_feat[:, 0].int().tolist() == expected
|
|
|
|
def test_legacy_fallback_multigroup_emits_warning(self, monkeypatch) -> None:
|
|
"""group_detr > 1 legacy checkpoint (no num_queries/group_detr in args) emits scramble-risk warning."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
num_queries=2,
|
|
num_select=2,
|
|
group_detr=3,
|
|
)
|
|
labelled_refpoint = _labelled_query_tensor(num_queries=4, group_detr=3, dim=4)
|
|
labelled_query_feat = _labelled_query_tensor(num_queries=4, group_detr=3, dim=256)
|
|
checkpoint = {
|
|
"model": {
|
|
"class_embed.weight": torch.randn(91, 256),
|
|
"class_embed.bias": torch.randn(91),
|
|
"refpoint_embed.weight": labelled_refpoint,
|
|
"query_feat.weight": labelled_query_feat,
|
|
},
|
|
"args": {}, # no num_queries / group_detr keys
|
|
}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
captured: list[str] = []
|
|
|
|
def _capture(msg: str, *args: object, **kwargs: object) -> None:
|
|
try:
|
|
captured.append(msg % args if args else msg)
|
|
except TypeError:
|
|
captured.append(msg)
|
|
|
|
monkeypatch.setattr("rfdetr.models.weights.logger.warning", _capture)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
assert any("group_detr" in msg and ("scramble" in msg or "flat slice" in msg) for msg in captured), (
|
|
f"Expected scramble-risk warning for group_detr > 1; got: {captured}"
|
|
)
|
|
|
|
def test_legacy_fallback_when_args_missing_num_queries_key(self, monkeypatch, tmp_path):
|
|
"""When checkpoint args dict lacks num_queries/group_detr keys, falls back to flat legacy slice."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mc = RFDETRBaseConfig(
|
|
pretrain_weights="/fake/weights.pth",
|
|
device="cpu",
|
|
num_queries=2,
|
|
num_select=2,
|
|
group_detr=1,
|
|
)
|
|
labelled_refpoint = _labelled_query_tensor(num_queries=4, group_detr=1, dim=4)
|
|
labelled_query_feat = _labelled_query_tensor(num_queries=4, group_detr=1, dim=256)
|
|
checkpoint = {
|
|
"model": {
|
|
"class_embed.weight": torch.randn(91, 256),
|
|
"class_embed.bias": torch.randn(91),
|
|
"refpoint_embed.weight": labelled_refpoint,
|
|
"query_feat.weight": labelled_query_feat,
|
|
},
|
|
"args": {},
|
|
}
|
|
monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
|
|
|
|
nn_model = _fake_nn_model()
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
passed_state = nn_model.load_state_dict.call_args[0][0]
|
|
refpoint = passed_state["refpoint_embed.weight"]
|
|
query_feat = passed_state["query_feat.weight"]
|
|
expected = [0, 1]
|
|
assert refpoint[:, 0].int().tolist() == expected
|
|
assert query_feat[:, 0].int().tolist() == expected
|
|
|
|
|
|
# Partial-load detector
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPartialLoadDetector:
|
|
"""Tests for ``_warn_on_partial_load`` — surfaces silent partial loads.
|
|
|
|
The rf-detr logger has ``propagate=False`` so pytest's ``caplog`` does not see its records. These tests monkeypatch
|
|
``logger.warning`` directly to capture the message text.
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def captured(self, monkeypatch):
|
|
"""Capture every call to ``rfdetr.models.weights.logger.warning`` as a formatted string."""
|
|
captured: list[str] = []
|
|
|
|
def _capture(msg, *args, **kwargs):
|
|
try:
|
|
captured.append(msg % args if args else msg)
|
|
except TypeError:
|
|
captured.append(msg)
|
|
|
|
monkeypatch.setattr("rfdetr.models.weights.logger.warning", _capture)
|
|
return captured
|
|
|
|
@pytest.mark.parametrize(
|
|
"result",
|
|
[
|
|
pytest.param(
|
|
SimpleNamespace(missing_keys=[], unexpected_keys=[]),
|
|
id="clean_load",
|
|
),
|
|
pytest.param(
|
|
SimpleNamespace(
|
|
missing_keys=[
|
|
"class_embed.weight",
|
|
"bbox_embed.layers.0.weight",
|
|
"refpoint_embed.weight",
|
|
"query_feat.weight",
|
|
"transformer.enc_out_class_embed.0.weight",
|
|
"transformer.enc_out_bbox_embed.0.layers.0.weight",
|
|
],
|
|
unexpected_keys=[],
|
|
),
|
|
id="intentional_head_keys",
|
|
),
|
|
pytest.param(
|
|
SimpleNamespace(missing_keys=42, unexpected_keys=[]),
|
|
id="non_iterable_missing_keys",
|
|
),
|
|
],
|
|
)
|
|
def test_no_warning_cases(self, captured, result: SimpleNamespace) -> None:
|
|
"""Cases that must not emit any partial-load warning.
|
|
|
|
Covers: clean load, intentional head keys, and non-iterable missing_keys.
|
|
"""
|
|
_warn_on_partial_load(result, "/fake/weights.pth")
|
|
assert captured == []
|
|
|
|
def test_unexpected_backbone_missing_keys_warn(self, captured):
|
|
"""Missing backbone keys (e.g. register_tokens) must trigger the warning."""
|
|
result = SimpleNamespace(
|
|
missing_keys=[
|
|
"backbone.0.encoder.encoder.embeddings.register_tokens",
|
|
"backbone.0.encoder.encoder.layers.0.register_block.weight",
|
|
],
|
|
unexpected_keys=[],
|
|
)
|
|
_warn_on_partial_load(result, "/fake/weights.pth")
|
|
assert len(captured) == 1
|
|
assert "/fake/weights.pth" in captured[0]
|
|
assert "register_tokens" in captured[0]
|
|
|
|
def test_unexpected_keys_warn(self, captured):
|
|
"""Unexpected checkpoint keys (model has no slot for them) must trigger the warning."""
|
|
result = SimpleNamespace(
|
|
missing_keys=[],
|
|
unexpected_keys=["backbone.0.encoder.legacy_module.weight"],
|
|
)
|
|
_warn_on_partial_load(result, "/fake/weights.pth")
|
|
assert len(captured) == 1
|
|
assert "not consumed by model" in captured[0]
|
|
|
|
def test_removed_keypoint_projection_keys_do_not_warn(self, captured):
|
|
"""Legacy keypoint projection tensors are intentionally ignored during partial-load checks."""
|
|
result = SimpleNamespace(
|
|
missing_keys=[],
|
|
unexpected_keys=[
|
|
"keypoint_head.keypoint_proj.0.weight",
|
|
"keypoint_head.keypoint_proj.0.bias",
|
|
"keypoint_head.keypoint_proj.2.weight",
|
|
"keypoint_head.keypoint_proj.2.bias",
|
|
],
|
|
)
|
|
_warn_on_partial_load(result, "/fake/weights.pth")
|
|
assert captured == []
|
|
|
|
def test_removed_keypoint_projection_keys_do_not_mask_other_unexpected_keys(self, captured):
|
|
"""Only the removed keypoint projection tensors are filtered from the partial-load warning."""
|
|
result = SimpleNamespace(
|
|
missing_keys=[],
|
|
unexpected_keys=[
|
|
"keypoint_head.keypoint_proj.0.weight",
|
|
"keypoint_head.keypoint_proj.0.bias",
|
|
"keypoint_head.keypoint_proj.2.weight",
|
|
"keypoint_head.keypoint_proj.2.bias",
|
|
"backbone.0.encoder.legacy_module.weight",
|
|
],
|
|
)
|
|
_warn_on_partial_load(result, "/fake/weights.pth")
|
|
assert len(captured) == 1
|
|
assert "legacy_module" in captured[0]
|
|
assert "keypoint_proj" not in captured[0]
|
|
|
|
def test_handles_non_iterable_input_gracefully(self, captured):
|
|
"""A MagicMock-style result (used in many existing tests) must not raise."""
|
|
_warn_on_partial_load(MagicMock(), "/fake/weights.pth")
|
|
# The crucial assertion is "did not raise"; whether captured is empty
|
|
# depends on MagicMock truthiness — both outcomes are acceptable.
|
|
|
|
@pytest.mark.parametrize(
|
|
"missing_keys, unexpected_keys, count_str",
|
|
[
|
|
pytest.param(
|
|
[f"backbone.0.encoder.layer.{i}.weight" for i in range(10)],
|
|
[],
|
|
"10 model parameter",
|
|
id="long_missing_keys",
|
|
),
|
|
pytest.param(
|
|
[],
|
|
[f"backbone.0.legacy.{i}.weight" for i in range(8)],
|
|
"8 checkpoint key(s)",
|
|
id="long_unexpected_keys",
|
|
),
|
|
],
|
|
)
|
|
def test_truncates_long_key_lists_in_message(
|
|
self,
|
|
captured,
|
|
missing_keys: list[str],
|
|
unexpected_keys: list[str],
|
|
count_str: str,
|
|
) -> None:
|
|
"""Sample key lists in the warning are bounded to 5 entries with a trailing ellipsis."""
|
|
result = SimpleNamespace(missing_keys=missing_keys, unexpected_keys=unexpected_keys)
|
|
_warn_on_partial_load(result, "/fake/weights.pth")
|
|
assert len(captured) == 1
|
|
assert count_str in captured[0]
|
|
assert "..." in captured[0]
|
|
|
|
def test_mixed_intentional_and_unintentional_keys_warn_only_for_unexpected(self, captured) -> None:
|
|
"""Only unintentional missing keys appear in the warning; intentional reinit keys are filtered.
|
|
|
|
When a checkpoint load returns both head-reinit keys (class_embed.weight, etc.) and a genuine backbone mismatch
|
|
(backbone.0.encoder.register_tokens), the warning must fire exactly once and must reference the unintentional
|
|
key, not the filtered ones.
|
|
"""
|
|
result = SimpleNamespace(
|
|
missing_keys=[
|
|
"class_embed.weight",
|
|
"bbox_embed.layers.0.weight",
|
|
"refpoint_embed.weight",
|
|
"backbone.0.encoder.encoder.embeddings.register_tokens",
|
|
],
|
|
unexpected_keys=[],
|
|
)
|
|
_warn_on_partial_load(result, "/fake/mixed.pth")
|
|
assert len(captured) == 1, f"Expected exactly one warning, got {len(captured)}: {captured}"
|
|
assert "register_tokens" in captured[0], "Warning must reference the unintentional backbone key"
|
|
assert "class_embed" not in captured[0], "Intentional head key must be filtered from warning text"
|
|
|
|
@patch("rfdetr.models.weights.torch.load")
|
|
@patch("rfdetr.models.weights.os.path.isfile", return_value=True)
|
|
@patch("rfdetr.models.weights.validate_checkpoint_compatibility")
|
|
@patch("rfdetr.models.weights.validate_pretrain_weights")
|
|
@patch("rfdetr.models.weights.download_pretrain_weights")
|
|
def test_partial_load_is_invoked_during_load_pretrain_weights(
|
|
self,
|
|
mock_download: MagicMock,
|
|
mock_validate_weights: MagicMock,
|
|
mock_validate_compat: MagicMock,
|
|
mock_isfile: MagicMock,
|
|
mock_torch_load: MagicMock,
|
|
monkeypatch,
|
|
) -> None:
|
|
"""Integration check: load_pretrain_weights wires up the partial-load detector."""
|
|
from rfdetr.models.weights import load_pretrain_weights
|
|
|
|
mock_torch_load.return_value = _make_checkpoint(num_classes=91)
|
|
|
|
mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
|
|
nn_model = _fake_nn_model()
|
|
nn_model.load_state_dict.return_value = SimpleNamespace(
|
|
missing_keys=["backbone.0.encoder.something_required.weight"],
|
|
unexpected_keys=[],
|
|
)
|
|
|
|
captured: list[str] = []
|
|
|
|
def _capture_warning(msg, *args, **kwargs):
|
|
try:
|
|
captured.append(msg % args if args else msg)
|
|
except TypeError:
|
|
captured.append(msg)
|
|
|
|
monkeypatch.setattr("rfdetr.models.weights.logger.warning", _capture_warning)
|
|
load_pretrain_weights(nn_model, mc)
|
|
|
|
assert any("partially" in m for m in captured), (
|
|
f"Expected partial-load warning to fire; got messages: {captured}"
|
|
)
|
|
|
|
# Conversely, a clean load must not fire the warning.
|
|
captured.clear()
|
|
nn_model.load_state_dict.return_value = SimpleNamespace(missing_keys=[], unexpected_keys=[])
|
|
load_pretrain_weights(nn_model, mc)
|
|
assert not any("partially" in m for m in captured), f"Clean load must not warn; got messages: {captured}"
|