File size: 2,836 Bytes
e6f80f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """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: # masked layout training: a given request is a condition, never a target
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"]
# The implemented BEV operator deliberately retains the frozen complete
# geometry-label policy. Yaw is always a learned head; it is not fixed by
# a request field, even when position or size coordinates are prescribed.
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):
# All ranks enter this boundary, including ranks with zero local terms.
count = torch.tensor(self.local_count, dtype=torch.int64, device=accelerator.device)
return int(accelerator.reduce(count, reduction="sum").item())
|