| """Label/term eligibility shared by training preflight and the criterion. |
| |
| Eligibility is independent of numeric loss: a correctly predicted supervised |
| example remains eligible. Scene regularizers need reliable input geometry, |
| not fabricated target labels. Counts describe eligible terms, not coordinates. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
|
|
|
|
| def supervision_masks(batch, validity=None): |
| validity = batch["validity"] if validity is None else validity |
| mask = batch["slot_mask"] |
| if "given_pose_mask" in batch: |
| mask = mask & ~batch["given_pose_mask"] |
| fields = {} |
| for key in ("position", "size"): |
| learn = ~batch.get(f"fixed_{key}_mask", torch.zeros_like(validity[key])) |
| fields[key] = mask & validity[key].all(-1) & learn.any(-1) |
| fields["yaw"] = mask & validity["yaw"] |
| |
| |
| |
| fields["box"] = mask & validity["position"].all(-1) & validity["size"].all(-1) & validity["yaw"] |
| return fields |
|
|
|
|
| def local_objective_count(batch, config, validity=None): |
| masks = supervision_masks(batch, validity) |
| count = sum(int(masks[key].sum()) for key in ("position", "size", "box") if getattr(config, key) > 0) |
| if config.yaw_cls > 0 or config.yaw_reg > 0: |
| count += int(masks["yaw"].sum()) |
| if config.collision or config.boundary: |
| from fastfill.v2.regularizers import regularizer_specs |
| if "conditions" not in batch: |
| raise ValueError("scene regularizers require original conditions") |
| for condition in batch["conditions"]: |
| objects, _, pairs, fixed_pairs = regularizer_specs(condition, config) |
| if config.boundary: |
| count += len(objects) |
| if config.collision: |
| count += len(pairs) + len(fixed_pairs) |
| return count |
|
|
|
|
| @dataclass(frozen=True) |
| class ObjectiveWindow: |
| """Accumulate local eligibility until Accelerate synchronizes a window.""" |
|
|
| local_count: int = 0 |
|
|
| def add(self, local_count): |
| if isinstance(local_count, bool) or not isinstance(local_count, int) or local_count < 0: |
| raise ValueError("active objective count must be a nonnegative integer") |
| return ObjectiveWindow(self.local_count + local_count) |
|
|
| def global_count(self, accelerator): |
| |
| count = torch.tensor(self.local_count, dtype=torch.int64, device=accelerator.device) |
| return int(accelerator.reduce(count, reduction="sum").item()) |
|
|
