Coverage for collectra/tasks/object_detection/rfdetr_feature_fusion.py: 88.18%
203 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Shared RGB DINOv2 encoding and attention pooling across aligned RGB views.
3Checkpoints use state dictionaries, never serialized Python model objects.
4The fusion specification is a persistent buffer, so upstream training/EMA
5checkpoints remain reconstructible even when custom top-level metadata is lost.
6"""
8from __future__ import annotations
10import contextlib
11from pathlib import Path
12import tempfile
13import types
15import torch
16from torch import nn
17from rfdetr.models.backbone.backbone import Backbone
18from rfdetr.models.backbone.dinov2 import DinoV2
20SPEC_KEY = "backbone.0.encoder.fusion_spec"
21FORMAT_VERSION = 1
24class ViewAttention(nn.Module):
25 """Reference-conditioned spatial attention over views, not over image pixels."""
27 def __init__(
28 self, channels: int, num_views: int, reference_view: int, hidden_dim: int
29 ):
30 super().__init__()
31 self.reference_view = reference_view
32 self.score = nn.Sequential(
33 nn.Linear(2 * channels, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1)
34 )
35 # Begin with uniform attention; the residual gate preserves RGB exactly.
36 nn.init.zeros_(self.score[-1].weight)
37 nn.init.zeros_(self.score[-1].bias)
38 self.view_bias = nn.Parameter(torch.zeros(num_views))
39 self.gate = nn.Parameter(torch.zeros(()))
41 def forward(self, views: torch.Tensor) -> torch.Tensor:
42 # B,V,C,H,W -> B,H,W,V,C. Attention is local to aligned positions.
43 features = views.permute(0, 3, 4, 1, 2)
44 reference = features[..., self.reference_view, :]
45 query = reference.unsqueeze(-2).expand_as(features)
46 logits = self.score(torch.cat((features, query), dim=-1)).squeeze(-1)
47 weights = (logits + self.view_bias).float().softmax(dim=-1).to(features.dtype)
48 pooled = (features * weights.unsqueeze(-1)).sum(dim=-2)
49 fused = reference + self.gate.tanh() * (pooled - reference)
50 return fused.permute(0, 3, 1, 2).contiguous()
53class MultiViewDinoV2(DinoV2):
54 def forward(self, pixels: torch.Tensor) -> list[torch.Tensor]:
55 if pixels.ndim != 4 or pixels.shape[1] != 3 * self.num_views:
56 raise ValueError(
57 f"Expected BCHW with {3 * self.num_views} channels; got {tuple(pixels.shape)}"
58 )
59 batch, _, height, width = pixels.shape
60 # One encoder invocation batches all views; parameters are shared.
61 rgb = pixels.reshape(batch * self.num_views, 3, height, width)
62 features = super().forward(rgb)
63 if len(features) != len(self.view_fusion):
64 raise RuntimeError(
65 "Encoder feature levels changed; cannot apply saved fusion modules"
66 )
67 return [
68 fusion(feature.reshape(batch, self.num_views, *feature.shape[1:]))
69 for feature, fusion in zip(features, self.view_fusion)
70 ]
72 def export(self) -> None:
73 raise NotImplementedError(
74 "Feature fusion currently supports eager prediction/training, not ONNX export"
75 )
78class FeatureFusionBackbone(Backbone):
79 def get_named_param_lr_pairs(self, args, prefix="backbone.0"):
80 pairs = super().get_named_param_lr_pairs(args, prefix)
81 # New fusion parameters must not inherit the tiny pretrained embedding LR.
82 for name, param in self.encoder.view_fusion.named_parameters():
83 pairs[f"{prefix}.encoder.view_fusion.{name}"] = {
84 "params": param,
85 "lr": args.lr,
86 "weight_decay": args.weight_decay,
87 }
88 return pairs
91def fusion_spec(state: dict) -> dict | None:
92 value = state.get(SPEC_KEY)
93 if value is None:
94 return None
95 # EMA can convert integer buffers to floats and introduce roundoff.
96 # Truncation would turn e.g. 24.999998 views into 24. The learned
97 # parameter shapes are authoritative for reconstruction, including older
98 # checkpoints whose specification has already been truncated.
99 version, views, reference, hidden = (round(v) for v in value.tolist())
100 prefix = "backbone.0.encoder.view_fusion.0."
101 if prefix + "view_bias" in state:
102 views = state[prefix + "view_bias"].shape[0]
103 if prefix + "score.0.weight" in state:
104 hidden = state[prefix + "score.0.weight"].shape[0]
105 if version != FORMAT_VERSION:
106 raise ValueError(f"Unsupported feature-fusion checkpoint version: {version}")
107 if views < 1 or not 0 <= reference < views or hidden < 1:
108 raise ValueError("Invalid feature-fusion checkpoint specification")
109 return {"num_views": views, "reference_view": reference, "hidden_dim": hidden}
112def install_feature_fusion(
113 nn_model, num_views: int, reference_view: int = 0, hidden_dim: int = 64
114) -> None:
115 if num_views < 1 or not 0 <= reference_view < num_views or hidden_dim < 1:
116 raise ValueError(
117 "Require num_views >= 1, 0 <= reference_view < num_views, hidden_dim >= 1"
118 )
119 backbone = nn_model.backbone[0]
120 encoder = backbone.encoder
121 if type(backbone) is not Backbone or type(encoder) is not DinoV2:
122 raise ValueError(
123 "Expected an unmodified RF-DETR DINOv2 backbone (LoRA/export/fusion wrappers are unsupported)"
124 )
125 projection = encoder.encoder.embeddings.patch_embeddings.projection
126 if projection.in_channels != 3:
127 raise ValueError(
128 "Feature fusion requires an RGB checkpoint, not an expanded input projection"
129 )
130 # Preserve every existing weight and key; only add fusion parameters/buffer.
131 encoder.__class__ = MultiViewDinoV2
132 backbone.__class__ = FeatureFusionBackbone
133 encoder.num_views = num_views
134 encoder.view_fusion = nn.ModuleList(
135 [
136 ViewAttention(channels, num_views, reference_view, hidden_dim)
137 for channels in encoder._out_feature_channels
138 ]
139 ).to(device=projection.weight.device, dtype=projection.weight.dtype)
140 encoder.register_buffer(
141 "fusion_spec",
142 torch.tensor(
143 [FORMAT_VERSION, num_views, reference_view, hidden_dim],
144 dtype=torch.int64,
145 device=projection.weight.device,
146 ),
147 )
148 encoder.view_fusion.train(encoder.training)
151def _sync(detector, spec: dict) -> None:
152 channels = 3 * spec["num_views"]
153 detector.model_config.num_channels = channels
154 detector.model.args.num_channels = channels
155 detector.means = [0.485, 0.456, 0.406] * spec["num_views"]
156 detector.stds = [0.229, 0.224, 0.225] * spec["num_views"]
157 detector.train = types.MethodType(_train_feature_fusion, detector)
160def adapt_model(
161 detector, num_views: int, reference_view: int = 0, hidden_dim: int = 64
162):
163 if detector.model_config.backbone_lora:
164 raise ValueError("LoRA backbones are not supported by feature fusion")
165 spec = dict(
166 num_views=num_views, reference_view=reference_view, hidden_dim=hidden_dim
167 )
168 install_feature_fusion(detector.model.model, **spec)
169 _sync(detector, spec)
170 return detector
173def save_adapted_model(detector, path: str | Path) -> None:
174 state = detector.model.model.state_dict()
175 spec = fusion_spec(state)
176 if spec is None:
177 raise ValueError("Model has no feature-fusion specification")
178 config = detector.model_config.model_dump(mode="json")
179 config["model_name"] = type(detector).__name__
180 args = dict(config, class_names=list(detector.model.class_names or []))
181 path = Path(path)
182 path.parent.mkdir(parents=True, exist_ok=True)
183 torch.save(
184 {
185 "model": state,
186 "model_name": type(detector).__name__,
187 "model_config": config,
188 "args": args,
189 "feature_fusion": spec,
190 },
191 path,
192 )
195def load_adapted_model(path: str | Path, device: str = "cpu"):
196 import rfdetr
198 checkpoint = torch.load(path, map_location="cpu", weights_only=True)
199 state = checkpoint["model"]
200 spec = fusion_spec(state)
201 if spec is None:
202 raise ValueError("Checkpoint is not an RF-DETR feature-fusion model")
203 args = checkpoint.get("args", {})
204 notes = args.get("notes") or {}
205 saved_config = checkpoint.get("model_config")
206 if saved_config is None and isinstance(notes, dict):
207 saved_config = notes.get("collectra_feature_fusion", {}).get("model_config")
208 if saved_config is None:
209 raise ValueError("Fusion checkpoint has no reconstruction config")
210 config = dict(saved_config)
211 name = checkpoint.get("model_name") or config.get("model_name")
212 if not isinstance(name, str) or not name.startswith("RFDETR"):
213 raise ValueError("Checkpoint has no RFDETR model class")
214 # Reconstruct RGB architecture before installing fusion. Never tile input weights.
215 config.update(pretrain_weights=None, num_channels=3, device=device)
216 config["num_classes"] = state["class_embed.weight"].shape[0] - 1
217 detector = getattr(rfdetr, name)(**config)
218 adapt_model(detector, **spec)
219 detector.model.model.load_state_dict(state, strict=True)
220 detector.model.class_names = checkpoint.get("args", {}).get("class_names", [])
221 detector.model_config.pretrain_weights = str(Path(path).resolve())
222 return detector
225@contextlib.contextmanager
226def feature_fusion_training(spec: dict):
227 """Install fusion before upstream loads weights and creates the optimizer."""
228 import rfdetr.training.module_model as module_model
230 original = module_model.build_model_from_config
232 def build(model_config, *args, **kwargs):
233 model = original(model_config, *args, **kwargs)
234 install_feature_fusion(model, **spec)
235 # Upstream strips model_config from best_total. TrainConfig.notes lives
236 # inside args and survives stripping, including custom resolutions.
237 train_config = args[0] if args else kwargs.get("train_config")
238 if train_config is not None:
239 previous_notes = train_config.notes
240 notes = (
241 dict(previous_notes)
242 if isinstance(previous_notes, dict)
243 else (
244 {"user_notes": previous_notes} if previous_notes is not None else {}
245 )
246 )
247 notes["collectra_feature_fusion"] = {
248 "model_config": model_config.model_dump(mode="json"),
249 **spec,
250 }
251 train_config.notes = notes
252 return model
254 module_model.build_model_from_config = build
255 try:
256 yield
257 finally:
258 module_model.build_model_from_config = original
261def _train_feature_fusion(detector, **kwargs):
262 """Train the current fusion weights rather than an unmodified RGB rebuild."""
263 from rfdetr.detr import RFDETR
264 from .rfdetr_data import multichannel_datasets
265 from .rfdetr_channels import require_multichannel_rfdetr
267 require_multichannel_rfdetr()
268 spec = fusion_spec(detector.model.model.state_dict())
269 if spec is None:
270 raise ValueError("Missing feature-fusion specification")
271 # Runtime model-builder patches do not propagate to independently spawned ranks.
272 devices = kwargs.get("devices", 1)
273 strategy = kwargs.get("strategy", "auto")
274 if (
275 devices not in (1, "1")
276 or kwargs.get("num_nodes", 1) != 1
277 or strategy not in ("auto", "single_device")
278 ):
279 raise ValueError(
280 "Feature fusion currently supports single-device training only"
281 )
282 augmentation = kwargs.pop("augmentation", "default")
283 if augmentation not in ("default", "none"):
284 raise ValueError(
285 "Feature fusion supports default/none torchvision augmentation"
286 )
287 if kwargs.get("aug_config") not in (None, {}):
288 raise ValueError(
289 "Nonempty Albumentations presets cannot preserve multichannel TIFFs"
290 )
291 if augmentation == "none":
292 kwargs["aug_config"] = {}
293 kwargs["augmentation_backend"] = "torchvision"
294 kwargs.setdefault("devices", 1)
295 kwargs.setdefault("save_dataset_grids", False)
296 previous = detector.model_config.pretrain_weights
297 # RF-DETR constructs a separate training module. Snapshot current weights,
298 # including learned attention, so repeated training calls retain them.
299 with tempfile.TemporaryDirectory(prefix="rfdetr-feature-fusion-") as folder:
300 snapshot = Path(folder) / "initial.pth"
301 save_adapted_model(detector, snapshot)
302 detector.model_config.pretrain_weights = str(snapshot)
303 try:
304 with feature_fusion_training(spec), multichannel_datasets():
305 RFDETR.train(detector, **kwargs)
306 finally:
307 detector.model_config.pretrain_weights = previous
310def expand_rfdetr_feature_fusion(
311 input_model: Path,
312 output_model: Path,
313 num_views: int,
314 variant: str | None = None,
315 reference_view: int = 0,
316 hidden_dim: int = 64,
317):
318 import rfdetr
320 if variant is None:
321 detector = rfdetr.RFDETR.from_checkpoint(str(input_model), device="cpu")
322 else:
323 names = {
324 "nano": "RFDETRNano",
325 "small": "RFDETRSmall",
326 "medium": "RFDETRMedium",
327 "base": "RFDETRBase",
328 "large": "RFDETRLarge",
329 "xlarge": "RFDETRXLarge",
330 "2xlarge": "RFDETR2XLarge",
331 }
332 name = names.get(variant.lower(), variant)
333 if not name.startswith("RFDETR") or not hasattr(rfdetr, name):
334 raise ValueError(f"Unknown RF-DETR variant: {variant}")
335 detector = getattr(rfdetr, name)(
336 pretrain_weights=str(input_model), device="cpu"
337 )
338 adapt_model(detector, num_views, reference_view, hidden_dim)
339 save_adapted_model(detector, output_model)
340 return detector
342def validate_saved_model(path: Path, expected_views: int) -> None:
343 import torch
345 detector = load_adapted_model(path)
346 encoder = detector.model.model.backbone[0].encoder
347 if encoder.num_views != expected_views:
348 raise RuntimeError("Reloaded view count does not match")
349 detector.model.model.eval()
350 # Exercise the real encoder/attention/projector rather than only inspecting weights.
351 block = detector.model_config.patch_size * detector.model_config.num_windows
352 with torch.no_grad():
353 features = encoder(torch.zeros(1, 3 * expected_views, block, block))
354 projected = detector.model.model.backbone[0].projector(features)
355 if not projected or any(not torch.isfinite(feature).all() for feature in projected):
356 raise RuntimeError("Feature-fusion forward produced invalid features")