Coverage for collectra/tasks/image_segmentation/sam3_metrics.py: 94.19%
155 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Optional SAM3 training metrics, reporting and weights-only checkpoint hooks."""
3import csv
4import json
5import logging
6import math
7import os
8import tempfile
9from pathlib import Path
12class MaskIoU:
13 """Pixel-weighted foreground IoU across validation queries and all ranks."""
15 def __init__(self, confidence_threshold=0.5):
16 self.confidence_threshold = confidence_threshold
17 self.reset()
19 def reset(self):
20 self.intersection = self.union = self.queries = 0
22 def update(self, *, find_stages, batch, **kwargs):
23 import torch
24 from torch.nn import functional as F
26 with torch.no_grad():
27 for outputs, target in zip(find_stages, batch.find_targets):
28 masks = outputs["pred_masks"].detach()
29 scores = (
30 outputs["pred_logits"].detach().sigmoid().reshape(masks.shape[:2])
31 )
32 presence = outputs.get("presence_logit_dec")
33 if presence is not None:
34 scores = scores * presence.detach().sigmoid().reshape(-1, 1)
35 offset = 0
36 for query, count in enumerate(target.num_boxes.tolist()):
37 segments = target.segments[offset : offset + count]
38 validity = target.is_valid_segment
39 valid = validity is None or bool(
40 validity[offset : offset + count].all()
41 )
42 offset += count
43 if not valid:
44 continue
45 truth = (
46 segments.bool().any(dim=0)
47 if count
48 else torch.zeros(
49 batch.img_batch.shape[-2:],
50 device=masks.device,
51 dtype=torch.bool,
52 )
53 )
54 predicted = torch.zeros_like(truth)
55 selected = masks[query, scores[query] >= self.confidence_threshold]
56 # Bound temporary memory when many instances are predicted.
57 for chunk in selected.split(8):
58 if len(chunk):
59 pixels = (
60 F.interpolate(
61 chunk[:, None].float(),
62 size=truth.shape,
63 mode="bilinear",
64 align_corners=False,
65 )[:, 0]
66 > 0
67 )
68 predicted |= pixels.any(dim=0)
69 self.intersection += int((predicted & truth).sum())
70 self.union += int((predicted | truth).sum())
71 self.queries += 1
73 def compute(self):
74 """Local progress value; never synchronize from a progress callback."""
75 return {
76 "iou": self.intersection / self.union
77 if self.union
78 else (1.0 if self.queries else float("nan"))
79 }
81 def compute_synced(self):
82 import torch
83 import torch.distributed as dist
85 distributed = dist.is_available() and dist.is_initialized()
86 device = (
87 torch.device("cuda", torch.cuda.current_device())
88 if distributed and dist.get_backend() == "nccl"
89 else "cpu"
90 )
91 totals = torch.tensor(
92 [self.intersection, self.union, self.queries],
93 dtype=torch.float64,
94 device=device,
95 )
96 if distributed:
97 dist.all_reduce(totals)
98 intersection, union, queries = totals.tolist()
99 return {
100 "iou": intersection / union if union else (1.0 if queries else float("nan"))
101 }
104class EpochReporter:
105 """Keep one CSV row per epoch, merging train and validation reports."""
107 def __init__(self, output, *, wandb=False, run_name=None, project=None, config=None):
108 self.output = Path(output)
109 self.output.mkdir(parents=True, exist_ok=True)
110 self.rows = {}
111 self.run = None
112 if wandb:
113 import wandb as sdk
115 self.run = sdk.init(
116 project=project or "collectra",
117 name=run_name,
118 dir=str(self.output),
119 config=config,
120 )
121 self.run.define_metric("epoch")
122 self.run.define_metric("*", step_metric="epoch")
124 def record(self, epoch, phase, metrics):
125 values = {}
126 for key, value in metrics.items():
127 if hasattr(value, "item"):
128 value = value.item()
129 if isinstance(value, (int, float)) and math.isfinite(value):
130 values[f"{phase}/{key}"] = value
131 row = self.rows.setdefault(epoch, {"epoch": epoch})
132 row.update(values)
133 columns = ["epoch"] + sorted(
134 {key for row in self.rows.values() for key in row if key != "epoch"}
135 )
136 temporary = self.output / "metrics.csv.tmp"
137 with temporary.open("w", newline="", encoding="utf-8") as stream:
138 writer = csv.DictWriter(stream, fieldnames=columns)
139 writer.writeheader()
140 writer.writerows(self.rows[key] for key in sorted(self.rows))
141 os.replace(temporary, self.output / "metrics.csv")
142 if self.run is not None:
143 self.run.log({"epoch": epoch, **values})
145 def close(self):
146 if self.run is not None:
147 self.run.finish()
150class TrainingHooks:
151 """Mixin for Meta's trainer; model/backend imports remain lazy."""
153 def __init__(
154 self,
155 *,
156 collectra_output,
157 collectra_wandb=False,
158 collectra_run_name=None,
159 collectra_project=None,
160 **kwargs,
161 ):
162 self._collectra_best = None
163 self._collectra_output = Path(collectra_output)
164 self._collectra_has_validation = bool(kwargs["data"].get("val"))
165 self._collectra_reporter = None
166 super().__init__(**kwargs)
167 if self.distributed_rank == 0:
168 self._collectra_reporter = EpochReporter(
169 collectra_output,
170 wandb=collectra_wandb,
171 run_name=collectra_run_name,
172 project=collectra_project,
173 config={
174 "epochs": self.max_epochs,
175 "selection_metric": (
176 "val/mask_iou"
177 if self._collectra_has_validation
178 else "train/loss"
179 ),
180 },
181 )
183 def load_checkpoint(self):
184 # Each invocation fine-tunes the explicitly configured model weights.
185 # Old optimizer checkpoints in a reused output directory must not resume.
186 self._init_model_state()
188 def save_checkpoint(self, epoch, checkpoint_names=None):
189 # Suppress native periodic, last-epoch and per-meter optimizer snapshots.
190 # Selection happens only after the epoch's metric has been computed.
191 return
193 def _save_best(self, score, metric, *, maximize):
194 import torch
196 if not math.isfinite(score):
197 raise ValueError(f"Cannot select a best SAM3 model: {metric} is not finite")
198 improved = self._collectra_best is None or (
199 score > self._collectra_best if maximize else score < self._collectra_best
200 )
201 if not improved:
202 return
203 self._collectra_best = score
204 if self.distributed_rank != 0:
205 return
206 directory = self._collectra_output / "training/checkpoints"
207 directory.mkdir(parents=True, exist_ok=True)
208 model = (
209 self.model.module
210 if isinstance(self.model, torch.nn.parallel.DistributedDataParallel)
211 else self.model
212 )
213 descriptor, name = tempfile.mkstemp(
214 prefix=".best-", suffix=".pt", dir=directory
215 )
216 os.close(descriptor)
217 try:
218 torch.save(model.state_dict(), name)
219 os.replace(name, directory / "best.pt")
220 finally:
221 Path(name).unlink(missing_ok=True)
222 metadata = {"epoch": self.epoch + 1, "metric": metric, "value": score}
223 (directory / "best.json").write_text(json.dumps(metadata, indent=2) + "\n")
224 logging.info(
225 "Best SAM3 weights: epoch=%s, %s=%s", self.epoch + 1, metric, score
226 )
228 def train_epoch(self, train_loader):
229 metrics = super().train_epoch(train_loader)
230 loss = metrics.get("Losses/train_all_loss")
231 if loss is not None:
232 metrics["loss"] = float(loss)
233 if not self._collectra_has_validation:
234 if loss is None:
235 raise ValueError("SAM3 training did not report its training loss")
236 self._save_best(float(loss), "train/loss", maximize=False)
237 if self._collectra_reporter is not None:
238 self._collectra_reporter.record(self.epoch + 1, "train", metrics)
239 return metrics
241 def val_epoch(self, val_loader, phase):
242 metrics = super().val_epoch(val_loader, phase)
243 values = [
244 value
245 for key, value in metrics.items()
246 if key.replace("\\", "/").endswith("/mask_iou/iou")
247 ]
248 if len(values) != 1:
249 raise ValueError("SAM3 validation did not report its mask IoU")
250 metrics["mask_iou"] = float(values[0])
251 loss = metrics.get("Losses/val_roboflow100_loss")
252 if loss is not None:
253 metrics["loss"] = float(loss)
254 self._save_best(metrics["mask_iou"], "val/mask_iou", maximize=True)
255 if self._collectra_reporter is not None:
256 self._collectra_reporter.record(self.epoch + 1, "val", metrics)
257 return metrics
259 def run(self):
260 try:
261 return super().run()
262 finally:
263 if self._collectra_reporter is not None:
264 self._collectra_reporter.close()
267def create_trainer(**kwargs):
268 from sam3.train.trainer import Trainer
270 class CollectraTrainer(TrainingHooks, Trainer):
271 pass
273 return CollectraTrainer(**kwargs)