Coverage for collectra/tasks/object_detection/yolo_feature_fusion.py: 85.58%
215 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Shared RGB YOLO backbone with attention fusion before the neck.
3Custom model classes live in this importable module so native Ultralytics .pt
4checkpoints survive a new Python process. Use load_adapted_model (or Collectra)
5for training; ordinary YOLO reconstruction deliberately fails on the tagged
6YAML instead of silently throwing away attention weights.
7"""
9from __future__ import annotations
11from copy import deepcopy
12from pathlib import Path
14import torch
15from torch import nn
16from ultralytics import YOLO
17from ultralytics.models.yolo.detect import DetectionTrainer
18from ultralytics.nn.tasks import DetectionModel
20FORMAT_VERSION = 1
21RGB_CONV_TAG = "CollectraFeatureFusionRGBConv"
24class ViewAttention(nn.Module):
25 """Reference-conditioned, per-position softmax attention over RGB views."""
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.Conv2d(2 * channels, hidden_dim, 1),
34 nn.GELU(),
35 nn.Conv2d(hidden_dim, 1, 1),
36 )
37 nn.init.zeros_(self.score[-1].weight)
38 nn.init.zeros_(self.score[-1].bias)
39 self.view_bias = nn.Parameter(torch.zeros(num_views))
40 self.gate = nn.Parameter(torch.zeros(()))
42 def forward(self, features: torch.Tensor) -> torch.Tensor:
43 batch, views, channels, height, width = features.shape
44 reference = features[:, self.reference_view]
45 query = reference[:, None].expand_as(features)
46 paired = torch.cat((features, query), dim=2).reshape(
47 batch * views, 2 * channels, height, width
48 )
49 logits = self.score(paired).reshape(batch, views, height, width)
50 weights = (
51 (logits + self.view_bias[None, :, None, None])
52 .float()
53 .softmax(1)
54 .to(features.dtype)
55 )
56 pooled = (features * weights[:, :, None]).sum(1)
57 return reference + self.gate.tanh() * (pooled - reference)
60def _validate_spec(spec: dict) -> None:
61 if spec.get("version") != FORMAT_VERSION:
62 raise ValueError("Unsupported YOLO feature-fusion checkpoint version")
63 views, reference, hidden = (
64 spec["num_views"],
65 spec["reference_view"],
66 spec["hidden_dim"],
67 )
68 if views < 1 or not 0 <= reference < views or hidden < 1:
69 raise ValueError(
70 "Require num_views >= 1, 0 <= reference_view < num_views, hidden_dim >= 1"
71 )
74def _layer_input(layer, current, saved):
75 if layer.f == -1:
76 return current
77 return (
78 saved[layer.f]
79 if isinstance(layer.f, int)
80 else [current if index == -1 else saved[index] for index in layer.f]
81 )
84class FeatureFusionDetectionModel(DetectionModel):
85 """Batched per-view backbone followed by fused multiscale neck inputs."""
87 def __init__(self, cfg: dict, ch: int | None = None, nc=None, verbose=True):
88 config = deepcopy(cfg)
89 spec = config.pop("feature_fusion", None)
90 if spec is None:
91 raise ValueError("A feature-fusion model configuration is required")
92 _validate_spec(spec)
93 if ch is not None and ch != 3 * spec["num_views"]:
94 raise ValueError(
95 f"Dataset has {ch} channels; fusion model expects {3 * spec['num_views']}"
96 )
97 if config["backbone"][0][2] == RGB_CONV_TAG:
98 config["backbone"][0][2] = "Conv"
99 config["channels"] = 3
100 # Stride initialization in DetectionModel calls _predict_once before
101 # fusion is installed; its RGB path remains available during construction.
102 super().__init__(cfg=config, ch=3, nc=nc, verbose=False)
103 self._install_fusion(spec)
104 if verbose:
105 self.info()
107 def _install_fusion(self, spec: dict) -> None:
108 _validate_spec(spec)
109 if hasattr(self, "view_fusion"):
110 raise ValueError("Model is already adapted for feature fusion")
111 first = self.model[0]
112 if (
113 self.yaml["backbone"][0][2] != "Conv"
114 or not hasattr(first, "conv")
115 or first.conv.in_channels != 3
116 ):
117 raise ValueError(
118 "Expected an unmodified RGB YOLO detection backbone starting with Conv"
119 )
120 if not hasattr(first, "bn"):
121 raise ValueError(
122 "Use an unfused training checkpoint, not a Conv/BatchNorm-fused inference model"
123 )
124 self.num_views = spec["num_views"]
125 self.backbone_layers = len(self.yaml["backbone"])
126 # Fuse every saved backbone output plus its final output. This preserves
127 # the full graph, including all backbone-to-neck skip connections.
128 indices = sorted(
129 {index for index in self.save if index < self.backbone_layers}
130 | {self.backbone_layers - 1}
131 )
132 parameter = next(self.parameters())
133 flags = [(module, module.training) for module in self.modules()]
134 self.eval()
135 try:
136 size = max(64, 2 * int(self.stride.max()))
137 current = torch.zeros(
138 1, 3, size, size, device=parameter.device, dtype=parameter.dtype
139 )
140 saved, channels = [], {}
141 with torch.no_grad():
142 for layer in self.model[: self.backbone_layers]:
143 current = layer(_layer_input(layer, current, saved))
144 if not isinstance(current, torch.Tensor) or current.ndim != 4:
145 raise ValueError(
146 "Only tensor feature maps from standard YOLO detection backbones are supported"
147 )
148 saved.append(current if layer.i in self.save else None)
149 if layer.i in indices:
150 channels[str(layer.i)] = current.shape[1]
151 finally:
152 for module, training in flags:
153 module.training = training
154 self.view_fusion = nn.ModuleDict(
155 {
156 index: ViewAttention(
157 width, self.num_views, spec["reference_view"], spec["hidden_dim"]
158 )
159 for index, width in channels.items()
160 }
161 ).to(device=parameter.device, dtype=parameter.dtype)
162 self.view_fusion.train(self.training)
163 self.yaml = deepcopy(self.yaml)
164 self.yaml["channels"] = 3 * self.num_views
165 self.yaml["feature_fusion"] = dict(spec)
166 # A vanilla trainer must not rebuild a widened RGB model and silently
167 # discard fusion. Our trainer removes this marker before RGB construction.
168 self.yaml["backbone"][0][2] = RGB_CONV_TAG
169 self.task = "detect"
171 def _fuse_feature(
172 self, index: int, feature: torch.Tensor, batch: int
173 ) -> torch.Tensor:
174 return self.view_fusion[str(index)](
175 feature.reshape(batch, self.num_views, *feature.shape[1:])
176 )
178 def _predict_once(self, x, profile=False, embed=None):
179 if not hasattr(self, "view_fusion"):
180 return super()._predict_once(x, profile=profile, embed=embed)
181 if profile or embed:
182 raise NotImplementedError(
183 "Feature-fusion layer profiling/embedding extraction is not supported"
184 )
185 if x.ndim != 4 or x.shape[1] != 3 * self.num_views:
186 raise ValueError(
187 f"Expected BCHW with {3 * self.num_views} channels; got {tuple(x.shape)}"
188 )
189 batch, _, height, width = x.shape
190 current = x.reshape(batch * self.num_views, 3, height, width)
191 saved = []
192 for layer in self.model:
193 if layer.i == self.backbone_layers:
194 last = self.backbone_layers - 1
195 for index in self.view_fusion:
196 i = int(index)
197 if saved[i] is not None:
198 saved[i] = self._fuse_feature(i, saved[i], batch)
199 current = (
200 saved[last]
201 if saved[last] is not None
202 else self._fuse_feature(last, current, batch)
203 )
204 current = layer(_layer_input(layer, current, saved))
205 saved.append(current if layer.i in self.save else None)
206 return current
208 def load(self, weights, verbose=True):
209 source = (
210 (weights.get("ema") or weights["model"])
211 if isinstance(weights, dict)
212 else weights
213 )
214 if not isinstance(source, FeatureFusionDetectionModel):
215 raise ValueError(
216 "Load an RGB checkpoint through adapt_model first; training requires fusion weights"
217 )
218 if source.yaml["feature_fusion"] != self.yaml["feature_fusion"]:
219 raise ValueError("Source and destination fusion configurations differ")
220 # Class-count changes are handled by Ultralytics; fusion weights must all transfer.
221 expected = {
222 key: value
223 for key, value in self.state_dict().items()
224 if key.startswith("view_fusion.")
225 }
226 state = source.state_dict()
227 if any(
228 key not in state or state[key].shape != value.shape
229 for key, value in expected.items()
230 ):
231 raise ValueError("Checkpoint has missing or incompatible attention weights")
232 return super().load(weights, verbose=verbose)
235class FeatureFusionDetectionTrainer(DetectionTrainer):
236 def get_model(self, cfg=None, weights=None, verbose=True):
237 if not isinstance(cfg, dict) or "feature_fusion" not in cfg:
238 raise ValueError("Feature-fusion training requires an adapted checkpoint")
239 model = self.set_model_names_for_load(
240 FeatureFusionDetectionModel(
241 cfg, ch=self.data["channels"], nc=self.data["nc"], verbose=verbose
242 )
243 )
244 if weights is not None:
245 model.load(weights, verbose=verbose)
246 return model
248 def plot_training_samples(self, batch, ni):
249 """RGB training plots cannot represent concatenated RGB views."""
252class FeatureFusionYOLO(YOLO):
253 @property
254 def task_map(self):
255 mapping = super().task_map
256 mapping["detect"]["model"] = FeatureFusionDetectionModel
257 mapping["detect"]["trainer"] = FeatureFusionDetectionTrainer
258 return mapping
260 def train(self, trainer=None, **kwargs):
261 if trainer is not None and trainer is not FeatureFusionDetectionTrainer:
262 raise ValueError(
263 "Use FeatureFusionDetectionTrainer to preserve the fusion architecture"
264 )
265 device = kwargs.get("device", self.overrides.get("device"))
266 if isinstance(device, (list, tuple)) or (
267 isinstance(device, str) and "," in device
268 ):
269 raise ValueError(
270 "Feature fusion currently supports single-device training only"
271 )
272 first = self.model.model[0]
273 if not hasattr(first, "bn"):
274 raise ValueError(
275 "Prediction fused BatchNorm; reload the unfused checkpoint before training"
276 )
277 kwargs.setdefault("plots", False)
278 return super().train(trainer=FeatureFusionDetectionTrainer, **kwargs)
280 def export(self, **kwargs):
281 raise NotImplementedError(
282 "YOLO feature fusion currently supports native PyTorch, not exported runtimes"
283 )
286def enable_feature_fusion(yolo):
287 """Upgrade a loaded YOLO wrapper so training reconstructs its custom model."""
288 if isinstance(yolo.model, FeatureFusionDetectionModel):
289 _validate_spec(yolo.model.yaml["feature_fusion"])
290 if type(yolo) is YOLO:
291 yolo.__class__ = FeatureFusionYOLO
292 elif not isinstance(yolo, FeatureFusionYOLO):
293 raise ValueError("Only standard YOLO detection wrappers are supported")
294 return yolo
297def adapt_model(yolo, num_views: int, reference_view: int = 0, hidden_dim: int = 64):
298 spec = dict(
299 version=FORMAT_VERSION,
300 num_views=num_views,
301 reference_view=reference_view,
302 hidden_dim=hidden_dim,
303 )
304 _validate_spec(spec)
305 if (
306 yolo.task != "detect"
307 or type(yolo.model) is not DetectionModel
308 or type(yolo) is not YOLO
309 ):
310 raise ValueError(
311 "Expected a standard, unmodified Ultralytics YOLO RGB detection model"
312 )
313 yolo.model.__class__ = FeatureFusionDetectionModel
314 yolo.model._install_fusion(spec)
315 return enable_feature_fusion(yolo)
318def save_adapted_model(yolo, path: str | Path) -> None:
319 if not isinstance(yolo.model, FeatureFusionDetectionModel):
320 raise ValueError("Model is not adapted for YOLO feature fusion")
321 path = Path(path)
322 if path.suffix.lower() != ".pt":
323 raise ValueError(
324 "Native Ultralytics feature-fusion checkpoints must use a .pt suffix"
325 )
326 path.parent.mkdir(parents=True, exist_ok=True)
327 previous = yolo.ckpt
328 try:
329 yolo.ckpt = previous or {}
330 yolo.save(str(path))
331 finally:
332 yolo.ckpt = previous
335def load_adapted_model(path: str | Path):
336 yolo = YOLO(str(path))
337 if not isinstance(yolo.model, FeatureFusionDetectionModel):
338 raise ValueError("Checkpoint is not a YOLO feature-fusion model")
339 return enable_feature_fusion(yolo)
342def expand_yolo_feature_fusion(
343 input_model: Path,
344 output_model: Path,
345 num_views: int,
346 reference_view: int = 0,
347 hidden_dim: int = 64,
348):
349 from ultralytics import YOLO
351 yolo = YOLO(str(input_model))
352 adapt_model(yolo, num_views, reference_view, hidden_dim)
353 save_adapted_model(yolo, output_model)
354 return yolo
356def validate_saved_model(path: Path, expected_views: int) -> None:
357 yolo = load_adapted_model(path)
358 model = yolo.model.eval()
359 if (
360 model.num_views != expected_views
361 or model.yaml["channels"] != 3 * expected_views
362 ):
363 raise RuntimeError("Reloaded model has inconsistent input channels")
364 if model.model[0].conv.in_channels != 3:
365 raise RuntimeError("Shared RGB backbone projection was unexpectedly widened")
366 size = max(64, 2 * int(model.stride.max()))
367 parameter = next(model.parameters())
368 with torch.no_grad():
369 result = model(
370 torch.zeros(
371 1,
372 3 * expected_views,
373 size,
374 size,
375 device=parameter.device,
376 dtype=parameter.dtype,
377 )
378 )
380 def tensors(value):
381 if isinstance(value, torch.Tensor):
382 yield value
383 elif isinstance(value, dict):
384 for item in value.values():
385 yield from tensors(item)
386 elif isinstance(value, (tuple, list)):
387 for item in value:
388 yield from tensors(item)
390 predictions = list(tensors(result))
391 if not predictions or any(not torch.isfinite(value).all() for value in predictions):
392 raise RuntimeError("Feature-fusion forward produced invalid predictions")