Add floor-event timing objective and addressee diagnostics
This commit is contained in:
Xin Wang
832
train_easycom_floor_event_probe.py
Normal file
832
train_easycom_floor_event_probe.py
Normal file
@@ -0,0 +1,832 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Train an analysis-only direct EasyCom response/yield event probe.
|
||||
|
||||
The probe consumes frozen Omni states, the same mixed-audio vector and one
|
||||
aggregate oracle-history vector at the causal decision source. Unlike the
|
||||
main activity head, it is supervised directly on competing-risk-censored
|
||||
response/yield labels. It is a diagnostic for objective and information
|
||||
alignment, not a deployable model or a replacement for dense activity mAP.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
import random
|
||||
from typing import Any, Iterable
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from easycom_head import (
|
||||
NUM_HORIZONS,
|
||||
STEP_SECONDS,
|
||||
TurnPredictionHead,
|
||||
load_session_data,
|
||||
)
|
||||
from evaluate_floor_state_changes import (
|
||||
DECISION_TIMINGS,
|
||||
OFFICIAL_TEST_SESSIONS,
|
||||
_rate_at_threshold,
|
||||
_safe_ap,
|
||||
_threshold_for_recall,
|
||||
collect_floor_event_candidates,
|
||||
normalize_session_name,
|
||||
)
|
||||
|
||||
|
||||
EVENT_NAMES = ("response", "yield")
|
||||
PROBE_ORACLE_CONDITIONS = (
|
||||
"aggregate-visible",
|
||||
"aggregate-visible-address-control",
|
||||
"aggregate-visible-addressee",
|
||||
)
|
||||
DEFAULT_TRAIN_SESSIONS = tuple(f"Session_{index}" for index in range(4, 12))
|
||||
DEFAULT_DEV_SESSIONS = ("Session_12",)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProbeExample:
|
||||
session_index: int
|
||||
source_row: int
|
||||
event_index: int
|
||||
horizon_index: int
|
||||
label: bool
|
||||
short_resume_hard_negative: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProbeSession:
|
||||
name: str
|
||||
states: torch.Tensor
|
||||
audio: torch.Tensor
|
||||
oracle: torch.Tensor
|
||||
oracle_valid: torch.Tensor
|
||||
target: np.ndarray
|
||||
oracle_condition: str
|
||||
oracle_contract: dict[str, Any]
|
||||
|
||||
|
||||
class FloorEventDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
sessions: list[ProbeSession],
|
||||
*,
|
||||
decision_timing: str,
|
||||
midturn_pause_max_steps: int,
|
||||
) -> None:
|
||||
self.sessions = sessions
|
||||
self.examples: list[ProbeExample] = []
|
||||
self.censoring: dict[str, dict[str, dict[str, int]]] = {}
|
||||
for session_index, session in enumerate(sessions):
|
||||
self.censoring[session.name] = {}
|
||||
session_ids = np.zeros(len(session.target), dtype=np.int32)
|
||||
for event_index, event_name in enumerate(EVENT_NAMES):
|
||||
aggregate = {
|
||||
"num_source_boundaries": 0,
|
||||
"num_right_censored": 0,
|
||||
"num_simultaneous_boundary_outcomes_excluded": 0,
|
||||
"num_zero_window_competing_risk": 0,
|
||||
}
|
||||
for horizon_index in range(NUM_HORIZONS):
|
||||
candidates, diagnostics = collect_floor_event_candidates(
|
||||
session.target,
|
||||
session_ids,
|
||||
horizon_index,
|
||||
event_type=event_name, # type: ignore[arg-type]
|
||||
decision_timing=decision_timing,
|
||||
midturn_pause_max_steps=midturn_pause_max_steps,
|
||||
)
|
||||
for key, value in diagnostics.items():
|
||||
aggregate[key] += int(value)
|
||||
self.examples.extend(
|
||||
ProbeExample(
|
||||
session_index=session_index,
|
||||
source_row=item.source_row,
|
||||
event_index=event_index,
|
||||
horizon_index=horizon_index,
|
||||
label=item.label,
|
||||
short_resume_hard_negative=(
|
||||
item.short_resume_hard_negative
|
||||
),
|
||||
)
|
||||
for item in candidates
|
||||
)
|
||||
self.censoring[session.name][event_name] = aggregate
|
||||
if not self.examples:
|
||||
raise ValueError("No observable floor-event examples")
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.examples)
|
||||
|
||||
def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
|
||||
example = self.examples[index]
|
||||
session = self.sessions[example.session_index]
|
||||
row = example.source_row
|
||||
return (
|
||||
session.states[row].float(),
|
||||
session.audio[row].float(),
|
||||
session.oracle[row].float(),
|
||||
session.oracle_valid[row].float(),
|
||||
torch.tensor(example.event_index, dtype=torch.long),
|
||||
torch.tensor(example.horizon_index, dtype=torch.long),
|
||||
torch.tensor(float(example.label), dtype=torch.float32),
|
||||
torch.tensor(
|
||||
float(example.short_resume_hard_negative), dtype=torch.float32
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class DirectFloorEventHead(nn.Module):
|
||||
"""Event readout with the same fusion trunk as the dense oracle head."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
state_shape: tuple[int, ...],
|
||||
audio_dim: int,
|
||||
oracle_dim: int,
|
||||
layer_fusion: str,
|
||||
local_projection_dim: int,
|
||||
person_attention_dim: int,
|
||||
person_attention_heads: int,
|
||||
dropout: float,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if len(state_shape) not in {1, 2}:
|
||||
raise ValueError(f"Expected state shape [D] or [L,D], got {state_shape}")
|
||||
if layer_fusion not in {"fixed-mean", "concat-mlp"}:
|
||||
raise ValueError(f"Unknown layer fusion: {layer_fusion}")
|
||||
if min(
|
||||
local_projection_dim,
|
||||
person_attention_dim,
|
||||
person_attention_heads,
|
||||
audio_dim,
|
||||
oracle_dim,
|
||||
) <= 0:
|
||||
raise ValueError("All probe dimensions must be positive")
|
||||
self.state_shape = state_shape
|
||||
num_input_layers = 1 if len(state_shape) == 1 else state_shape[0]
|
||||
effective_layer_fusion = "single" if len(state_shape) == 1 else layer_fusion
|
||||
self.head = TurnPredictionHead(
|
||||
hidden_size=state_shape[-1],
|
||||
num_horizons=NUM_HORIZONS,
|
||||
num_classes=len(EVENT_NAMES),
|
||||
dropout=dropout,
|
||||
num_input_layers=num_input_layers,
|
||||
layer_fusion=effective_layer_fusion,
|
||||
local_visual_dim=oracle_dim,
|
||||
local_audio_dim=audio_dim,
|
||||
local_projection_dim=local_projection_dim,
|
||||
local_fusion="oracle-frame-set-attention",
|
||||
person_attention_dim=person_attention_dim,
|
||||
person_attention_heads=person_attention_heads,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
states: torch.Tensor,
|
||||
audio: torch.Tensor,
|
||||
oracle: torch.Tensor,
|
||||
oracle_valid: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
if oracle_valid.ndim == 1:
|
||||
oracle_valid = oracle_valid[:, None]
|
||||
logits = self.head(
|
||||
states,
|
||||
audio,
|
||||
oracle[:, None, :],
|
||||
oracle_valid.bool(),
|
||||
)
|
||||
return logits.transpose(1, 2)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--feature-root", type=Path, required=True)
|
||||
parser.add_argument("--target-root", type=Path, required=True)
|
||||
parser.add_argument("--oracle-feature-root", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", type=Path, required=True)
|
||||
parser.add_argument("--result-dir", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--train-sessions", nargs="+", default=list(DEFAULT_TRAIN_SESSIONS)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dev-sessions", nargs="+", default=list(DEFAULT_DEV_SESSIONS)
|
||||
)
|
||||
parser.add_argument("--eval-sessions", nargs="+", default=None)
|
||||
parser.add_argument("--eval-only", action="store_true")
|
||||
parser.add_argument("--epochs", type=int, default=20)
|
||||
parser.add_argument("--batch-size", type=int, default=256)
|
||||
parser.add_argument("--learning-rate", type=float, default=1e-4)
|
||||
parser.add_argument("--weight-decay", type=float, default=5e-5)
|
||||
parser.add_argument("--dropout", type=float, default=0.2)
|
||||
parser.add_argument("--local-projection-dim", type=int, default=256)
|
||||
parser.add_argument("--person-attention-dim", type=int, default=64)
|
||||
parser.add_argument("--person-attention-heads", type=int, default=4)
|
||||
parser.add_argument(
|
||||
"--layer-fusion", choices=("fixed-mean", "concat-mlp"), default="concat-mlp"
|
||||
)
|
||||
parser.add_argument("--max-positive-weight", type=float, default=20.0)
|
||||
parser.add_argument("--matched-recall", type=float, default=0.8)
|
||||
parser.add_argument("--midturn-pause-max-seconds", type=float, default=2.0)
|
||||
parser.add_argument(
|
||||
"--decision-timing", choices=DECISION_TIMINGS, default="pre-boundary"
|
||||
)
|
||||
parser.add_argument("--expected-mode", default="open-stream-replay")
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--num-workers", type=int, default=0)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def normalize_sessions(values: Iterable[str | int]) -> list[str]:
|
||||
sessions = [normalize_session_name(value) for value in values]
|
||||
if len(sessions) != len(set(sessions)):
|
||||
raise ValueError(f"Duplicate sessions: {sessions}")
|
||||
return sessions
|
||||
|
||||
|
||||
def load_probe_session(args: argparse.Namespace, session: str) -> ProbeSession:
|
||||
data = load_session_data(
|
||||
args.feature_root,
|
||||
args.target_root,
|
||||
session,
|
||||
expected_mode=None if args.expected_mode == "any" else args.expected_mode,
|
||||
dataset="easycom",
|
||||
local_feature_root=args.oracle_feature_root,
|
||||
local_modalities="av",
|
||||
local_fusion="oracle-frame-set-attention",
|
||||
)
|
||||
if data.local_features is None or data.local_person_features is None:
|
||||
raise AssertionError("Oracle loader did not return local features")
|
||||
people, mask = data.local_person_features.padded()
|
||||
if people.shape[1] > 1:
|
||||
raise ValueError(
|
||||
f"{session}: direct probe requires an aggregate oracle condition "
|
||||
f"with at most one row-level vector, got max_people={people.shape[1]}"
|
||||
)
|
||||
counts = mask.sum(dim=1, keepdim=True)
|
||||
oracle = (people * mask[..., None]).sum(dim=1)
|
||||
oracle = oracle / counts.clamp_min(1).to(oracle.dtype)
|
||||
oracle_valid = mask.any(dim=1)
|
||||
contract = data.feature_metadata.get("local_contract")
|
||||
if not isinstance(contract, dict):
|
||||
raise ValueError(f"{session}: missing validated oracle contract")
|
||||
config = contract.get("config")
|
||||
condition = config.get("oracle_condition") if isinstance(config, dict) else None
|
||||
if condition not in PROBE_ORACLE_CONDITIONS:
|
||||
raise ValueError(
|
||||
f"{session}: expected an aggregate-visible oracle condition, got {condition!r}"
|
||||
)
|
||||
history_steps = config.get("history_steps")
|
||||
if not isinstance(history_steps, int) or history_steps <= 0:
|
||||
raise ValueError(f"{session}: invalid oracle history_steps")
|
||||
expected_oracle_dim = (
|
||||
3 * history_steps
|
||||
if condition != "aggregate-visible"
|
||||
else 2 * history_steps
|
||||
)
|
||||
if oracle.shape[1] != expected_oracle_dim:
|
||||
raise ValueError(
|
||||
f"{session}: condition {condition} expects oracle_dim="
|
||||
f"{expected_oracle_dim}, got {oracle.shape[1]}"
|
||||
)
|
||||
if condition == "aggregate-visible-address-control" and bool(
|
||||
torch.count_nonzero(oracle[:, -history_steps:])
|
||||
):
|
||||
raise ValueError(f"{session}: addressee capacity-control channel is not zero")
|
||||
if condition == "aggregate-visible-addressee" and config.get(
|
||||
"uses_target_of_speech"
|
||||
) is not True:
|
||||
raise ValueError(f"{session}: addressee oracle lacks privileged-input flag")
|
||||
return ProbeSession(
|
||||
name=session,
|
||||
states=data.states,
|
||||
audio=data.local_features,
|
||||
oracle=oracle,
|
||||
oracle_valid=oracle_valid,
|
||||
target=data.targets.numpy(),
|
||||
oracle_condition=condition,
|
||||
oracle_contract=contract,
|
||||
)
|
||||
|
||||
|
||||
def validate_session_shapes(sessions: list[ProbeSession]) -> dict[str, Any]:
|
||||
if not sessions:
|
||||
raise ValueError("No sessions loaded")
|
||||
first = sessions[0]
|
||||
state_shape = tuple(int(value) for value in first.states.shape[1:])
|
||||
audio_width = int(first.audio.shape[1])
|
||||
oracle_dim = int(first.oracle.shape[1])
|
||||
condition = first.oracle_condition
|
||||
for session in sessions[1:]:
|
||||
observed = (
|
||||
tuple(int(value) for value in session.states.shape[1:]),
|
||||
int(session.audio.shape[1]),
|
||||
int(session.oracle.shape[1]),
|
||||
session.oracle_condition,
|
||||
)
|
||||
expected = (state_shape, audio_width, oracle_dim, condition)
|
||||
if observed != expected:
|
||||
raise ValueError(
|
||||
f"Feature contract changes at {session.name}: {observed} vs {expected}"
|
||||
)
|
||||
if session.oracle_contract != first.oracle_contract:
|
||||
raise ValueError(
|
||||
f"Oracle feature contract changes at {session.name}; "
|
||||
"regenerate sidecars under one configuration"
|
||||
)
|
||||
if audio_width < 2:
|
||||
raise ValueError("Expected [audio vector, valid bit]")
|
||||
return {
|
||||
"state_shape": state_shape,
|
||||
"audio_dim": audio_width - 1,
|
||||
"oracle_dim": oracle_dim,
|
||||
"oracle_condition": condition,
|
||||
"oracle_contract": first.oracle_contract,
|
||||
}
|
||||
|
||||
|
||||
def positive_weights(
|
||||
dataset: FloorEventDataset, max_weight: float
|
||||
) -> torch.Tensor:
|
||||
positive = torch.zeros((len(EVENT_NAMES), NUM_HORIZONS), dtype=torch.float64)
|
||||
total = torch.zeros_like(positive)
|
||||
for item in dataset.examples:
|
||||
total[item.event_index, item.horizon_index] += 1
|
||||
positive[item.event_index, item.horizon_index] += int(item.label)
|
||||
negative = total - positive
|
||||
weights = negative / positive.clamp_min(1)
|
||||
return weights.clamp(min=1.0, max=max_weight).float()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_dataset(
|
||||
model: DirectFloorEventHead,
|
||||
dataset: FloorEventDataset,
|
||||
*,
|
||||
batch_size: int,
|
||||
device: torch.device,
|
||||
num_workers: int,
|
||||
) -> dict[str, np.ndarray]:
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
num_workers=num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
)
|
||||
model.eval()
|
||||
scores: list[np.ndarray] = []
|
||||
events: list[np.ndarray] = []
|
||||
horizons: list[np.ndarray] = []
|
||||
labels: list[np.ndarray] = []
|
||||
short_resume: list[np.ndarray] = []
|
||||
for batch in loader:
|
||||
states, audio, oracle, oracle_valid, event, horizon, label, short = batch
|
||||
logits = model(
|
||||
states.to(device, non_blocking=True),
|
||||
audio.to(device, non_blocking=True),
|
||||
oracle.to(device, non_blocking=True),
|
||||
oracle_valid.to(device, non_blocking=True),
|
||||
)
|
||||
row = torch.arange(len(event), device=device)
|
||||
selected = logits[
|
||||
row,
|
||||
event.to(device, non_blocking=True),
|
||||
horizon.to(device, non_blocking=True),
|
||||
]
|
||||
scores.append(selected.sigmoid().cpu().numpy())
|
||||
events.append(event.numpy())
|
||||
horizons.append(horizon.numpy())
|
||||
labels.append(label.numpy().astype(bool))
|
||||
short_resume.append(short.numpy().astype(bool))
|
||||
return {
|
||||
"score": np.concatenate(scores),
|
||||
"event": np.concatenate(events),
|
||||
"horizon": np.concatenate(horizons),
|
||||
"label": np.concatenate(labels),
|
||||
"short_resume": np.concatenate(short_resume),
|
||||
}
|
||||
|
||||
|
||||
def metric_rows(
|
||||
predictions: dict[str, np.ndarray],
|
||||
*,
|
||||
matched_recall: float,
|
||||
fixed_thresholds: dict[tuple[int, int], float | None] | None = None,
|
||||
threshold_source: str,
|
||||
) -> tuple[list[dict[str, Any]], dict[tuple[int, int], float | None]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
thresholds: dict[tuple[int, int], float | None] = {}
|
||||
for event_index, event_name in enumerate(EVENT_NAMES):
|
||||
for horizon_index in range(NUM_HORIZONS):
|
||||
mask = (predictions["event"] == event_index) & (
|
||||
predictions["horizon"] == horizon_index
|
||||
)
|
||||
score = predictions["score"][mask]
|
||||
label = predictions["label"][mask]
|
||||
positive_scores = score[label]
|
||||
negative_scores = score[~label]
|
||||
key = (event_index, horizon_index)
|
||||
threshold = (
|
||||
_threshold_for_recall(positive_scores, matched_recall)
|
||||
if fixed_thresholds is None
|
||||
else fixed_thresholds[key]
|
||||
)
|
||||
thresholds[key] = threshold
|
||||
hard_negative = predictions["short_resume"][mask] & ~label
|
||||
rows.append(
|
||||
{
|
||||
"event_type": event_name,
|
||||
"horizon_index": horizon_index + 1,
|
||||
"nominal_seconds": (horizon_index + 1) * STEP_SECONDS,
|
||||
"num_positive": int(label.sum()),
|
||||
"num_negative": int((~label).sum()),
|
||||
"positive_rate": float(label.mean()) if len(label) else None,
|
||||
"average_precision": _safe_ap(label, score),
|
||||
"matched_recall": matched_recall,
|
||||
"threshold": threshold,
|
||||
"threshold_source": threshold_source,
|
||||
"achieved_positive_recall": _rate_at_threshold(
|
||||
positive_scores, threshold
|
||||
),
|
||||
"negative_trigger_rate": _rate_at_threshold(
|
||||
negative_scores, threshold
|
||||
),
|
||||
"short_resume_hard_negative_count": int(hard_negative.sum()),
|
||||
"short_resume_trigger_rate": _rate_at_threshold(
|
||||
score[hard_negative], threshold
|
||||
),
|
||||
}
|
||||
)
|
||||
return rows, thresholds
|
||||
|
||||
|
||||
def mean_event_ap(rows: list[dict[str, Any]]) -> float:
|
||||
values = [
|
||||
float(row["average_precision"])
|
||||
for row in rows
|
||||
if row["average_precision"] is not None
|
||||
]
|
||||
return float(np.mean(values)) if values else float("-inf")
|
||||
|
||||
|
||||
def train_model(
|
||||
args: argparse.Namespace,
|
||||
model: DirectFloorEventHead,
|
||||
train_dataset: FloorEventDataset,
|
||||
dev_dataset: FloorEventDataset,
|
||||
device: torch.device,
|
||||
) -> tuple[dict[str, torch.Tensor], float, int, list[dict[str, Any]]]:
|
||||
generator = torch.Generator().manual_seed(args.seed)
|
||||
loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=True,
|
||||
generator=generator,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
)
|
||||
weights = positive_weights(train_dataset, args.max_positive_weight).to(device)
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(),
|
||||
lr=args.learning_rate,
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
best_score = float("-inf")
|
||||
best_epoch = 0
|
||||
best_state: dict[str, torch.Tensor] | None = None
|
||||
history: list[dict[str, Any]] = []
|
||||
model.to(device)
|
||||
for epoch in range(1, args.epochs + 1):
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
total_examples = 0
|
||||
for batch in loader:
|
||||
states, audio, oracle, oracle_valid, event, horizon, label, _short = batch
|
||||
states = states.to(device, non_blocking=True)
|
||||
audio = audio.to(device, non_blocking=True)
|
||||
oracle = oracle.to(device, non_blocking=True)
|
||||
oracle_valid = oracle_valid.to(device, non_blocking=True)
|
||||
event = event.to(device, non_blocking=True)
|
||||
horizon = horizon.to(device, non_blocking=True)
|
||||
label = label.to(device, non_blocking=True)
|
||||
logits = model(states, audio, oracle, oracle_valid)
|
||||
batch_rows = torch.arange(len(event), device=device)
|
||||
selected = logits[batch_rows, event, horizon]
|
||||
loss = F.binary_cross_entropy_with_logits(
|
||||
selected,
|
||||
label,
|
||||
weight=torch.where(label.bool(), weights[event, horizon], 1.0),
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += float(loss.detach()) * len(event)
|
||||
total_examples += len(event)
|
||||
|
||||
dev_predictions = predict_dataset(
|
||||
model,
|
||||
dev_dataset,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
dev_rows, _ = metric_rows(
|
||||
dev_predictions,
|
||||
matched_recall=args.matched_recall,
|
||||
threshold_source="development_sessions",
|
||||
)
|
||||
score = mean_event_ap(dev_rows)
|
||||
record = {
|
||||
"epoch": epoch,
|
||||
"train_loss": total_loss / max(total_examples, 1),
|
||||
"dev_mean_event_ap": score,
|
||||
}
|
||||
history.append(record)
|
||||
print(
|
||||
f"epoch={epoch:02d} train_loss={record['train_loss']:.6f} "
|
||||
f"dev_mean_event_ap={score:.6f}"
|
||||
)
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_epoch = epoch
|
||||
best_state = {
|
||||
key: value.detach().cpu().clone()
|
||||
for key, value in model.state_dict().items()
|
||||
}
|
||||
if best_state is None:
|
||||
raise RuntimeError("Training produced no checkpoint")
|
||||
return best_state, best_score, best_epoch, history
|
||||
|
||||
|
||||
def write_metrics_csv(path: Path, rows: list[dict[str, Any]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", newline="") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.epochs <= 0 or args.batch_size <= 0:
|
||||
raise ValueError("--epochs and --batch-size must be positive")
|
||||
if not 0 < args.matched_recall <= 1:
|
||||
raise ValueError("--matched-recall must be in (0,1]")
|
||||
if args.midturn_pause_max_seconds <= 0:
|
||||
raise ValueError("--midturn-pause-max-seconds must be positive")
|
||||
if args.max_positive_weight < 1:
|
||||
raise ValueError("--max-positive-weight must be at least 1")
|
||||
train_sessions = normalize_sessions(args.train_sessions)
|
||||
dev_sessions = normalize_sessions(args.dev_sessions)
|
||||
eval_sessions = normalize_sessions(args.eval_sessions or dev_sessions)
|
||||
if not args.eval_only:
|
||||
overlap = set(train_sessions) & set(dev_sessions)
|
||||
if overlap:
|
||||
raise ValueError(f"Training/development sessions overlap: {sorted(overlap)}")
|
||||
if set(train_sessions + dev_sessions) & set(OFFICIAL_TEST_SESSIONS):
|
||||
raise ValueError("Official test Sessions 1--3 cannot train or select the probe")
|
||||
if args.checkpoint.exists() and not args.overwrite:
|
||||
raise FileExistsError(
|
||||
f"Refusing to overwrite checkpoint {args.checkpoint}; pass --overwrite"
|
||||
)
|
||||
elif not args.checkpoint.is_file():
|
||||
raise FileNotFoundError(args.checkpoint)
|
||||
summary_path = args.result_dir / "summary.json"
|
||||
if summary_path.exists() and not args.overwrite:
|
||||
raise FileExistsError(
|
||||
f"Refusing to overwrite result {summary_path}; pass --overwrite"
|
||||
)
|
||||
|
||||
random.seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
torch.manual_seed(args.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
device = torch.device(args.device)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(f"CUDA requested but unavailable: {args.device}")
|
||||
midturn_steps = int(
|
||||
math.ceil(args.midturn_pause_max_seconds / STEP_SECONDS)
|
||||
)
|
||||
|
||||
requested = list(
|
||||
dict.fromkeys(
|
||||
eval_sessions
|
||||
if args.eval_only
|
||||
else train_sessions + dev_sessions + eval_sessions
|
||||
)
|
||||
)
|
||||
loaded = {
|
||||
session: load_probe_session(args, session)
|
||||
for session in requested
|
||||
}
|
||||
feature_config = validate_session_shapes(list(loaded.values()))
|
||||
model_config = {
|
||||
"state_shape": feature_config["state_shape"],
|
||||
"audio_dim": feature_config["audio_dim"],
|
||||
"oracle_dim": feature_config["oracle_dim"],
|
||||
"layer_fusion": args.layer_fusion,
|
||||
"local_projection_dim": args.local_projection_dim,
|
||||
"person_attention_dim": args.person_attention_dim,
|
||||
"person_attention_heads": args.person_attention_heads,
|
||||
"dropout": args.dropout,
|
||||
}
|
||||
print(
|
||||
f"condition={feature_config['oracle_condition']} "
|
||||
f"state_shape={feature_config['state_shape']} "
|
||||
f"audio_dim={feature_config['audio_dim']} "
|
||||
f"oracle_dim={feature_config['oracle_dim']} "
|
||||
f"decision_timing={args.decision_timing}"
|
||||
)
|
||||
|
||||
if args.eval_only:
|
||||
checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
|
||||
if checkpoint.get("format") != "easycom-direct-floor-event-probe-v1":
|
||||
raise ValueError(f"Invalid probe checkpoint format: {args.checkpoint}")
|
||||
saved_recall = float(checkpoint.get("matched_recall", float("nan")))
|
||||
if not math.isclose(saved_recall, args.matched_recall):
|
||||
raise ValueError(
|
||||
f"Checkpoint matched recall is {saved_recall}, requested "
|
||||
f"{args.matched_recall}"
|
||||
)
|
||||
train_sessions = list(checkpoint.get("train_sessions", []))
|
||||
dev_sessions = list(checkpoint.get("dev_sessions", []))
|
||||
if not dev_sessions:
|
||||
raise ValueError("Probe checkpoint has no development sessions")
|
||||
saved_config = checkpoint.get("model_config")
|
||||
if not isinstance(saved_config, dict):
|
||||
raise ValueError(f"Invalid probe checkpoint: {args.checkpoint}")
|
||||
comparable_saved = dict(saved_config)
|
||||
comparable_saved["state_shape"] = tuple(comparable_saved["state_shape"])
|
||||
if comparable_saved != model_config:
|
||||
# Evaluation must follow the checkpoint architecture; only the
|
||||
# feature dimensions/condition must match current data.
|
||||
expected_features = {
|
||||
key: comparable_saved[key]
|
||||
for key in ("state_shape", "audio_dim", "oracle_dim")
|
||||
}
|
||||
observed_features = {
|
||||
key: model_config[key]
|
||||
for key in ("state_shape", "audio_dim", "oracle_dim")
|
||||
}
|
||||
if expected_features != observed_features:
|
||||
raise ValueError(
|
||||
f"Checkpoint/data feature mismatch: {expected_features} vs "
|
||||
f"{observed_features}"
|
||||
)
|
||||
model_config = comparable_saved
|
||||
if checkpoint.get("oracle_condition") != feature_config["oracle_condition"]:
|
||||
raise ValueError(
|
||||
"Checkpoint/data oracle condition mismatch: "
|
||||
f"{checkpoint.get('oracle_condition')!r} vs "
|
||||
f"{feature_config['oracle_condition']!r}"
|
||||
)
|
||||
if checkpoint.get("decision_timing") != args.decision_timing:
|
||||
raise ValueError("Checkpoint/data decision timing mismatch")
|
||||
model = DirectFloorEventHead(**model_config)
|
||||
model.load_state_dict(checkpoint["model_state"])
|
||||
best_epoch = int(checkpoint["best_epoch"])
|
||||
best_score = float(checkpoint["best_dev_mean_event_ap"])
|
||||
history = checkpoint.get("history", [])
|
||||
dev_rows = checkpoint["dev_metrics"]
|
||||
raw_thresholds = checkpoint["dev_thresholds"]
|
||||
thresholds = {
|
||||
tuple(int(part) for part in key.split(",")): value
|
||||
for key, value in raw_thresholds.items()
|
||||
}
|
||||
else:
|
||||
train_dataset = FloorEventDataset(
|
||||
[loaded[name] for name in train_sessions],
|
||||
decision_timing=args.decision_timing,
|
||||
midturn_pause_max_steps=midturn_steps,
|
||||
)
|
||||
dev_dataset = FloorEventDataset(
|
||||
[loaded[name] for name in dev_sessions],
|
||||
decision_timing=args.decision_timing,
|
||||
midturn_pause_max_steps=midturn_steps,
|
||||
)
|
||||
print(
|
||||
f"train_examples={len(train_dataset)} "
|
||||
f"dev_examples={len(dev_dataset)}"
|
||||
)
|
||||
model = DirectFloorEventHead(**model_config)
|
||||
best_state, best_score, best_epoch, history = train_model(
|
||||
args, model, train_dataset, dev_dataset, device
|
||||
)
|
||||
model.load_state_dict(best_state)
|
||||
model.to(device)
|
||||
dev_predictions = predict_dataset(
|
||||
model,
|
||||
dev_dataset,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
dev_rows, thresholds = metric_rows(
|
||||
dev_predictions,
|
||||
matched_recall=args.matched_recall,
|
||||
threshold_source="development_sessions:" + ",".join(dev_sessions),
|
||||
)
|
||||
args.checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
||||
checkpoint_payload = {
|
||||
"format": "easycom-direct-floor-event-probe-v1",
|
||||
"model_config": model_config,
|
||||
"model_state": best_state,
|
||||
"oracle_condition": feature_config["oracle_condition"],
|
||||
"oracle_contract": feature_config["oracle_contract"],
|
||||
"decision_timing": args.decision_timing,
|
||||
"train_sessions": train_sessions,
|
||||
"dev_sessions": dev_sessions,
|
||||
"matched_recall": args.matched_recall,
|
||||
"best_epoch": best_epoch,
|
||||
"best_dev_mean_event_ap": best_score,
|
||||
"history": history,
|
||||
"dev_metrics": dev_rows,
|
||||
"dev_thresholds": {
|
||||
f"{event},{horizon}": value
|
||||
for (event, horizon), value in thresholds.items()
|
||||
},
|
||||
"args": {
|
||||
key: str(value) if isinstance(value, Path) else value
|
||||
for key, value in vars(args).items()
|
||||
},
|
||||
}
|
||||
torch.save(checkpoint_payload, args.checkpoint)
|
||||
|
||||
eval_dataset = FloorEventDataset(
|
||||
[loaded[name] for name in eval_sessions],
|
||||
decision_timing=args.decision_timing,
|
||||
midturn_pause_max_steps=midturn_steps,
|
||||
)
|
||||
print(f"eval_sessions={','.join(eval_sessions)} eval_examples={len(eval_dataset)}")
|
||||
model.to(device)
|
||||
eval_predictions = predict_dataset(
|
||||
model,
|
||||
eval_dataset,
|
||||
batch_size=args.batch_size,
|
||||
device=device,
|
||||
num_workers=args.num_workers,
|
||||
)
|
||||
eval_rows, _ = metric_rows(
|
||||
eval_predictions,
|
||||
matched_recall=args.matched_recall,
|
||||
fixed_thresholds=thresholds,
|
||||
threshold_source="development_sessions:" + ",".join(dev_sessions),
|
||||
)
|
||||
payload = {
|
||||
"protocol": "easycom_direct_floor_event_probe_v1",
|
||||
"analysis_only": True,
|
||||
"interpretation": (
|
||||
"Direct competing-risk response/yield supervision over frozen Omni, "
|
||||
"mixed audio and privileged aggregate oracle history. This diagnoses "
|
||||
"objective/information alignment and is not a deployable system."
|
||||
),
|
||||
"decision_timing": args.decision_timing,
|
||||
"oracle_condition": feature_config["oracle_condition"],
|
||||
"train_sessions": train_sessions,
|
||||
"dev_sessions": dev_sessions,
|
||||
"eval_sessions": eval_sessions,
|
||||
"matched_recall": args.matched_recall,
|
||||
"checkpoint": str(args.checkpoint.resolve()),
|
||||
"best_epoch": best_epoch,
|
||||
"best_dev_mean_event_ap": best_score,
|
||||
"model_config": model_config,
|
||||
"history": history,
|
||||
"dev_metrics": dev_rows,
|
||||
"eval_metrics": eval_rows,
|
||||
"eval_censoring_audit": eval_dataset.censoring,
|
||||
"num_eval_examples": len(eval_dataset),
|
||||
}
|
||||
args.result_dir.mkdir(parents=True, exist_ok=True)
|
||||
summary_path.write_text(json.dumps(payload, indent=2) + "\n")
|
||||
write_metrics_csv(args.result_dir / "metrics.csv", eval_rows)
|
||||
print(
|
||||
f"saved condition={feature_config['oracle_condition']} "
|
||||
f"best_epoch={best_epoch} dev_mean_event_ap={best_score:.6f}"
|
||||
)
|
||||
for event_name in EVENT_NAMES:
|
||||
event_rows = [row for row in eval_rows if row["event_type"] == event_name]
|
||||
print(
|
||||
f"{event_name}: mean_AP={mean_event_ap(event_rows):.6f}, "
|
||||
f"2s_AP={event_rows[-1]['average_precision']}, "
|
||||
f"2s_error_rate@R={args.matched_recall:g}="
|
||||
f"{event_rows[-1]['negative_trigger_rate']}"
|
||||
)
|
||||
print(f"summary: {summary_path}")
|
||||
print(f"metrics: {args.result_dir / 'metrics.csv'}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user