521 lines
21 KiB
Python
521 lines
21 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Continue stage-one fusion with selective MiniROADA GRU/head fine-tuning.
|
||
|
|
|
||
|
|
Frozen RGB/audio features and causal Omni states are reused. Final evaluation
|
||
|
|
exports stage-two, immutable stage-one and original EgoSpeak predictions under
|
||
|
|
identical masks, plus dense and threshold-free raw endpoint AP. No semantic
|
||
|
|
annotations, dev checkpoint selection or test-side tuning are used.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import csv
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from tqdm.auto import tqdm
|
||
|
|
|
||
|
|
from easycom_head import (
|
||
|
|
CLASS_NAMES,
|
||
|
|
NUM_HORIZONS,
|
||
|
|
OFFICIAL_TEST_SESSIONS,
|
||
|
|
OFFICIAL_TRAIN_SESSIONS,
|
||
|
|
)
|
||
|
|
from egospeak_omni_finetune import (
|
||
|
|
STAGE_TWO_FORMAT,
|
||
|
|
EgoSpeakOmniFinetune,
|
||
|
|
base_model_config,
|
||
|
|
load_stage_one_checkpoint,
|
||
|
|
validate_training_sessions,
|
||
|
|
)
|
||
|
|
from egospeak_omni_residual import (
|
||
|
|
EgoSpeakOmniResidual,
|
||
|
|
EgoSpeakOmniWindowDataset,
|
||
|
|
align_source_predictions,
|
||
|
|
miniroad_normalized_loss,
|
||
|
|
sha256_file,
|
||
|
|
)
|
||
|
|
from evaluate_floor_state_changes import evaluate_horizon
|
||
|
|
from train_egospeak_omni_residual import (
|
||
|
|
autocast_context,
|
||
|
|
evaluate_predictions,
|
||
|
|
load_sessions,
|
||
|
|
make_loader,
|
||
|
|
normalize_sessions,
|
||
|
|
omni_contract_signature,
|
||
|
|
predict_session,
|
||
|
|
print_metrics,
|
||
|
|
resolve_device,
|
||
|
|
set_seed,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--egospeak-feature-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--omni-feature-root", type=Path, required=True)
|
||
|
|
parser.add_argument("--stage-one-checkpoint", type=Path, required=True)
|
||
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
||
|
|
parser.add_argument("--train-sessions", nargs="+", default=None)
|
||
|
|
parser.add_argument("--eval-sessions", nargs="+", default=None)
|
||
|
|
parser.add_argument("--expected-omni-mode", default="open-stream-replay")
|
||
|
|
parser.add_argument(
|
||
|
|
"--omni-layer-fusion", choices=("single", "last", "mean"), default="mean"
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--base-trainable",
|
||
|
|
choices=("none", "heads", "gru-heads"),
|
||
|
|
default="gru-heads",
|
||
|
|
help="none is a matched extra-epoch fusion-only control; layer1 always stays frozen.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--epochs", type=int, default=4, help="Additional stage-two epochs."
|
||
|
|
)
|
||
|
|
parser.add_argument("--batch-size", type=int, default=64)
|
||
|
|
parser.add_argument("--num-workers", type=int, default=4)
|
||
|
|
parser.add_argument(
|
||
|
|
"--window-size", type=int, default=None, help="Inherit stage-one window."
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--stride", type=int, default=None, help="Inherit stage-one sampling stride."
|
||
|
|
)
|
||
|
|
parser.add_argument("--fusion-lr", type=float, default=1e-4)
|
||
|
|
parser.add_argument("--base-lr", type=float, default=1e-5)
|
||
|
|
parser.add_argument("--weight-decay", type=float, default=0.05)
|
||
|
|
parser.add_argument("--grad-clip", type=float, default=1.0)
|
||
|
|
parser.add_argument("--seed", type=int, default=0)
|
||
|
|
parser.add_argument("--device", default="cuda")
|
||
|
|
parser.add_argument("--amp", action="store_true")
|
||
|
|
parser.add_argument("--eval-chunk-size", type=int, default=0)
|
||
|
|
parser.add_argument("--max-train-examples", type=int, default=None)
|
||
|
|
parser.add_argument("--no-progress", action="store_true")
|
||
|
|
return parser.parse_args(argv)
|
||
|
|
|
||
|
|
|
||
|
|
def validate_checkpoint_contract(
|
||
|
|
args: argparse.Namespace, payload: dict[str, Any]
|
||
|
|
) -> None:
|
||
|
|
training = payload["training_config"]
|
||
|
|
for key in ("omni_layer_fusion", "expected_omni_mode"):
|
||
|
|
if key not in training or getattr(args, key) != training[key]:
|
||
|
|
raise ValueError(f"Stage-two {key} must match stage-one extraction/readout")
|
||
|
|
if args.window_size is None:
|
||
|
|
args.window_size = int(training["window_size"])
|
||
|
|
if args.stride is None:
|
||
|
|
args.stride = int(training["stride"])
|
||
|
|
if args.window_size != training["window_size"] or args.stride != training["stride"]:
|
||
|
|
raise ValueError(
|
||
|
|
"Keep stage-one window_size/stride fixed for this continuation"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@torch.inference_mode()
|
||
|
|
def verify_stage_one_recovery(
|
||
|
|
model: EgoSpeakOmniFinetune,
|
||
|
|
reference: EgoSpeakOmniResidual,
|
||
|
|
dataset: EgoSpeakOmniWindowDataset,
|
||
|
|
device: torch.device,
|
||
|
|
) -> float:
|
||
|
|
model.eval()
|
||
|
|
reference.eval()
|
||
|
|
rgb, audio, omni, tail, _, _, source = dataset[0]
|
||
|
|
inputs = (
|
||
|
|
rgb[None].to(device),
|
||
|
|
audio[None].to(device),
|
||
|
|
omni[None].to(device),
|
||
|
|
tail[None].to(device),
|
||
|
|
)
|
||
|
|
actual, actual_base = model.forward_last(*inputs)
|
||
|
|
expected, expected_base = reference.forward_last(*inputs)
|
||
|
|
maximum = max(
|
||
|
|
float((actual - expected).abs().max()),
|
||
|
|
float((actual_base - expected_base).abs().max()),
|
||
|
|
)
|
||
|
|
if maximum != 0.0:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Stage-two initialization changed stage-one logits: {maximum}"
|
||
|
|
)
|
||
|
|
print("Initialization audit: stage-two logits == trained stage-one logits exactly")
|
||
|
|
print(
|
||
|
|
f"Causal sample: source_row={int(source)}, cutoff={(int(source) + 1) / 5:.2f}s; "
|
||
|
|
"input ends at source row, labels are source+1,...,source+10"
|
||
|
|
)
|
||
|
|
return maximum
|
||
|
|
|
||
|
|
|
||
|
|
def train_epoch(model, loader, optimizer, scaler, device, args) -> dict[str, Any]:
|
||
|
|
model.train()
|
||
|
|
total_loss = 0.0
|
||
|
|
num_examples = 0
|
||
|
|
first_batch_gradient_norms = {}
|
||
|
|
iterator = tqdm(
|
||
|
|
loader, desc="train stage2", disable=args.no_progress, dynamic_ncols=True
|
||
|
|
)
|
||
|
|
for rgb, audio, omni, tail, targets, _, _ in iterator:
|
||
|
|
rgb = rgb.to(device=device, dtype=torch.float32)
|
||
|
|
audio = audio.to(device=device, dtype=torch.float32)
|
||
|
|
omni = omni.to(device=device, dtype=torch.float32)
|
||
|
|
tail = tail.to(device)
|
||
|
|
targets = targets.to(device=device, dtype=torch.float32)
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
with autocast_context(device, args.amp):
|
||
|
|
logits, _ = model.forward_last(rgb, audio, omni, tail)
|
||
|
|
loss = miniroad_normalized_loss(logits.float(), targets)
|
||
|
|
if not torch.isfinite(loss):
|
||
|
|
raise RuntimeError("Nonfinite stage-two training loss")
|
||
|
|
scaler.scale(loss).backward()
|
||
|
|
scaler.unscale_(optimizer)
|
||
|
|
if not num_examples:
|
||
|
|
for group in optimizer.param_groups:
|
||
|
|
norms = [
|
||
|
|
p.grad.detach().norm()
|
||
|
|
for p in group["params"]
|
||
|
|
if p.grad is not None
|
||
|
|
]
|
||
|
|
norm = float(torch.stack(norms).norm()) if norms else 0.0
|
||
|
|
first_batch_gradient_norms[group["name"]] = norm
|
||
|
|
print(
|
||
|
|
f"First-batch gradient norms (before clipping): {first_batch_gradient_norms}"
|
||
|
|
)
|
||
|
|
torch.nn.utils.clip_grad_norm_(
|
||
|
|
model.trainable_parameters(),
|
||
|
|
args.grad_clip,
|
||
|
|
error_if_nonfinite=True,
|
||
|
|
)
|
||
|
|
scaler.step(optimizer)
|
||
|
|
scaler.update()
|
||
|
|
num_examples += len(rgb)
|
||
|
|
total_loss += float(loss.detach()) * len(rgb)
|
||
|
|
iterator.set_postfix_str(f"loss={total_loss / num_examples:.5f}")
|
||
|
|
return {
|
||
|
|
"loss": total_loss / num_examples,
|
||
|
|
"num_examples": num_examples,
|
||
|
|
"first_batch_gradient_norms": first_batch_gradient_norms,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def raw_endpoint_metrics(predictions, valid, targets) -> dict[str, Any]:
|
||
|
|
names = list(predictions)
|
||
|
|
target = np.concatenate([targets[name] for name in names])
|
||
|
|
prediction = np.concatenate([predictions[name] for name in names])
|
||
|
|
mask = np.concatenate([valid[name] for name in names])
|
||
|
|
session_ids = np.concatenate(
|
||
|
|
[np.full(len(targets[name]), i) for i, name in enumerate(names)]
|
||
|
|
)
|
||
|
|
horizons = []
|
||
|
|
for horizon in range(NUM_HORIZONS):
|
||
|
|
item = evaluate_horizon(target, prediction, mask, session_ids, horizon, [])
|
||
|
|
horizons.append(
|
||
|
|
{
|
||
|
|
"seconds": item["nominal_seconds"],
|
||
|
|
"onset": item["wearer_off_to_on_endpoint"],
|
||
|
|
"offset": item["wearer_on_to_off_endpoint"],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
means = {}
|
||
|
|
for family in ("onset", "offset"):
|
||
|
|
values = [h[family]["average_precision"] for h in horizons]
|
||
|
|
# Undefined horizons stay undefined rather than being silently dropped.
|
||
|
|
means[f"mean_{family}_ap"] = (
|
||
|
|
None if any(v is None for v in values) else float(np.mean(values))
|
||
|
|
)
|
||
|
|
return {
|
||
|
|
"protocol": "causal_valid_endpoint_state_change_ap",
|
||
|
|
"horizons": horizons,
|
||
|
|
**means,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_output(path: Path) -> Path:
|
||
|
|
path = path.expanduser().resolve()
|
||
|
|
if path.exists() and (not path.is_dir() or any(path.iterdir())):
|
||
|
|
raise FileExistsError(f"Refusing to overwrite output: {path}")
|
||
|
|
path.mkdir(parents=True, exist_ok=True)
|
||
|
|
for method in (
|
||
|
|
path,
|
||
|
|
path / "stage_one_reference",
|
||
|
|
path / "frozen_egospeak_reference",
|
||
|
|
):
|
||
|
|
for name in ("raw_eval", "pred", "valid"):
|
||
|
|
(method / name).mkdir(parents=True)
|
||
|
|
return path
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str] | None = None) -> None:
|
||
|
|
args = parse_args(argv)
|
||
|
|
if (
|
||
|
|
args.epochs <= 0
|
||
|
|
or args.batch_size <= 0
|
||
|
|
or args.eval_chunk_size < 0
|
||
|
|
or args.num_workers < 0
|
||
|
|
or args.grad_clip <= 0
|
||
|
|
or args.weight_decay < 0
|
||
|
|
or not all(
|
||
|
|
math.isfinite(value)
|
||
|
|
for value in (
|
||
|
|
args.fusion_lr,
|
||
|
|
args.base_lr,
|
||
|
|
args.grad_clip,
|
||
|
|
args.weight_decay,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
):
|
||
|
|
raise ValueError("Invalid stage-two epoch/batch/chunk/worker/gradient settings")
|
||
|
|
set_seed(args.seed)
|
||
|
|
device = resolve_device(args.device)
|
||
|
|
if args.amp and device.type != "cuda":
|
||
|
|
raise ValueError("--amp requires CUDA")
|
||
|
|
reference, parent = load_stage_one_checkpoint(args.stage_one_checkpoint)
|
||
|
|
validate_checkpoint_contract(args, parent)
|
||
|
|
train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS)
|
||
|
|
eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS)
|
||
|
|
validate_training_sessions(train_names)
|
||
|
|
train_sessions = load_sessions(args, train_names)
|
||
|
|
eval_sessions = load_sessions(args, eval_names)
|
||
|
|
sessions = (*train_sessions, *eval_sessions)
|
||
|
|
if {int(s.omni.shape[-1]) for s in sessions} != {reference.omni_dim}:
|
||
|
|
raise ValueError("Omni feature dimensions differ from stage-one checkpoint")
|
||
|
|
signatures = {omni_contract_signature(s) for s in sessions}
|
||
|
|
if len(signatures) != 1:
|
||
|
|
raise ValueError("Train/evaluation Omni extraction contracts differ")
|
||
|
|
# v1 saves readout/mode but not the entire extraction contract in its
|
||
|
|
# checkpoint. Verify sibling metrics provenance when it is available.
|
||
|
|
parent_metrics = (
|
||
|
|
args.stage_one_checkpoint.expanduser().resolve().parent / "metrics.json"
|
||
|
|
)
|
||
|
|
contract_verified = False
|
||
|
|
if parent_metrics.is_file():
|
||
|
|
metadata = json.loads(parent_metrics.read_text()).get("omni_sources", {})
|
||
|
|
for source in metadata.values():
|
||
|
|
config = source["config"]
|
||
|
|
signature = (
|
||
|
|
source["model_id"],
|
||
|
|
source["mode"],
|
||
|
|
config.get("memory_seconds"),
|
||
|
|
config.get("video_fps"),
|
||
|
|
config.get("readout_last_n_layers"),
|
||
|
|
config.get("readout_layer_fusion"),
|
||
|
|
tuple(source["state_shape_before_layer_fusion"][1:]),
|
||
|
|
)
|
||
|
|
if signature not in signatures:
|
||
|
|
raise ValueError(
|
||
|
|
"Current Omni extraction differs from stage-one metrics provenance"
|
||
|
|
)
|
||
|
|
contract_verified = bool(metadata)
|
||
|
|
if not contract_verified:
|
||
|
|
print(
|
||
|
|
"WARNING: stage-one full extraction provenance unavailable; reuse its exact Omni root."
|
||
|
|
)
|
||
|
|
|
||
|
|
model = EgoSpeakOmniFinetune.from_stage_one(
|
||
|
|
reference,
|
||
|
|
parent["model_config"],
|
||
|
|
base_trainable=args.base_trainable,
|
||
|
|
).to(device)
|
||
|
|
reference = reference.to(device)
|
||
|
|
groups = model.optimizer_groups(fusion_lr=args.fusion_lr, base_lr=args.base_lr)
|
||
|
|
for group in groups:
|
||
|
|
print(
|
||
|
|
f"parameters {group['name']}: {sum(p.numel() for p in group['params']):,}; lr={group['lr']:g}"
|
||
|
|
)
|
||
|
|
print(f"base_trainable={args.base_trainable}; layer1 frozen=True; Omni frozen=True")
|
||
|
|
print(f"train={train_names}; final evaluation={eval_names}")
|
||
|
|
dataset = EgoSpeakOmniWindowDataset(
|
||
|
|
train_sessions,
|
||
|
|
window_size=args.window_size,
|
||
|
|
stride=args.stride,
|
||
|
|
seed=args.seed,
|
||
|
|
max_examples=args.max_train_examples,
|
||
|
|
)
|
||
|
|
difference = verify_stage_one_recovery(model, reference, dataset, device)
|
||
|
|
output = prepare_output(args.output_dir)
|
||
|
|
optimizer = torch.optim.AdamW(groups, weight_decay=args.weight_decay)
|
||
|
|
try:
|
||
|
|
scaler = torch.amp.GradScaler("cuda", enabled=args.amp)
|
||
|
|
except (AttributeError, TypeError): # Older PyTorch compatibility.
|
||
|
|
scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
|
||
|
|
history = []
|
||
|
|
for epoch in range(1, args.epochs + 1):
|
||
|
|
dataset.resample(epoch - 1)
|
||
|
|
result = train_epoch(
|
||
|
|
model, make_loader(dataset, args, epoch), optimizer, scaler, device, args
|
||
|
|
)
|
||
|
|
result["epoch"] = epoch
|
||
|
|
history.append(result)
|
||
|
|
print(
|
||
|
|
f"stage2 epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}; examples={result['num_examples']}"
|
||
|
|
)
|
||
|
|
|
||
|
|
training_config = vars(args).copy()
|
||
|
|
training_config = {
|
||
|
|
k: str(v) if isinstance(v, Path) else v for k, v in training_config.items()
|
||
|
|
}
|
||
|
|
budgets = {
|
||
|
|
"stage_one_epochs": int(parent["training_config"]["epochs"]),
|
||
|
|
"stage_two_epochs": args.epochs,
|
||
|
|
"total_adaptation_epochs": int(parent["training_config"]["epochs"])
|
||
|
|
+ args.epochs,
|
||
|
|
}
|
||
|
|
checkpoint = output / "checkpoint_last.pt"
|
||
|
|
provenance = {
|
||
|
|
"stage_one_checkpoint": str(args.stage_one_checkpoint.expanduser().resolve()),
|
||
|
|
"stage_one_checkpoint_sha256": sha256_file(
|
||
|
|
args.stage_one_checkpoint.expanduser().resolve()
|
||
|
|
),
|
||
|
|
"original_base_checkpoint": parent.get("base_checkpoint"),
|
||
|
|
"original_base_checkpoint_sha256": parent.get("base_checkpoint_sha256"),
|
||
|
|
"stage_one_full_extraction_contract_verified": contract_verified,
|
||
|
|
}
|
||
|
|
(output / "README.md").write_text(
|
||
|
|
"# Selective EgoSpeak + Omni fine-tuning\n\n"
|
||
|
|
f"Train: {train_names}. Final evaluation: {eval_names}. Seed: {args.seed}.\n\n"
|
||
|
|
f"Checkpoint: `checkpoint_last.pt`. Base policy: `{args.base_trainable}`. "
|
||
|
|
"RGB/audio feature extractors, input projection and Omni remain frozen.\n\n"
|
||
|
|
f"Epoch budgets: {budgets}. This continuation is not budget-matched "
|
||
|
|
"against stage one; use the `--base-trainable none` continuation control.\n\n"
|
||
|
|
"`metrics.json` reports dense AP/mAP and threshold-free raw endpoint AP "
|
||
|
|
"as fractions. `raw_endpoint_metrics.csv` has per-horizon AP and counts. "
|
||
|
|
"Raw onset uses OFF->ON versus OFF->OFF, scored by wearer probability; "
|
||
|
|
"offset uses ON->OFF versus ON->ON, scored by 1-wearer probability.\n\n"
|
||
|
|
"Immutable reference predictions are in `stage_one_reference/` and "
|
||
|
|
"`frozen_egospeak_reference/`, each with pred/valid/raw_eval and metrics. "
|
||
|
|
"No provisional semantic annotations are used.\n\n"
|
||
|
|
f"Run arguments:\n```json\n{json.dumps(training_config, indent=2)}\n```\n"
|
||
|
|
)
|
||
|
|
torch.save(
|
||
|
|
{
|
||
|
|
"format": STAGE_TWO_FORMAT,
|
||
|
|
"model_state_dict": model.state_dict(),
|
||
|
|
"base_model_config": base_model_config(model.base),
|
||
|
|
"model_config": parent["model_config"],
|
||
|
|
"base_trainable": args.base_trainable,
|
||
|
|
"training_config": training_config,
|
||
|
|
"train_sessions": train_names,
|
||
|
|
"eval_sessions": eval_names,
|
||
|
|
"stage_one_train_sessions": parent["train_sessions"],
|
||
|
|
"budgets": budgets,
|
||
|
|
"provenance": provenance,
|
||
|
|
"omni_sources": {item.name: item.omni_metadata for item in sessions},
|
||
|
|
"history": history,
|
||
|
|
"initialization_max_abs_stage_one_logit_difference": difference,
|
||
|
|
},
|
||
|
|
checkpoint,
|
||
|
|
)
|
||
|
|
|
||
|
|
directories = {
|
||
|
|
"stage_two": output,
|
||
|
|
"stage_one": output / "stage_one_reference",
|
||
|
|
"egospeak": output / "frozen_egospeak_reference",
|
||
|
|
}
|
||
|
|
predictions = {name: {} for name in directories}
|
||
|
|
valid = {}
|
||
|
|
targets = {}
|
||
|
|
for session in eval_sessions:
|
||
|
|
kwargs = {
|
||
|
|
"chunk_size": args.eval_chunk_size,
|
||
|
|
"amp": args.amp,
|
||
|
|
"show_progress": not args.no_progress,
|
||
|
|
}
|
||
|
|
fused, _ = predict_session(model, session, device, **kwargs)
|
||
|
|
first_stage, original = predict_session(reference, session, device, **kwargs)
|
||
|
|
for method, raw in (
|
||
|
|
("stage_two", fused),
|
||
|
|
("stage_one", first_stage),
|
||
|
|
("egospeak", original),
|
||
|
|
):
|
||
|
|
aligned, mask = align_source_predictions(raw, session.num_steps)
|
||
|
|
if session.name in valid and not np.array_equal(valid[session.name], mask):
|
||
|
|
raise RuntimeError("Reference and adapted valid masks differ")
|
||
|
|
valid[session.name] = mask
|
||
|
|
predictions[method][session.name] = aligned
|
||
|
|
directory = directories[method]
|
||
|
|
np.save(directory / "raw_eval" / f"{session.name}.npy", raw)
|
||
|
|
np.save(directory / "pred" / f"{session.name}.npy", aligned)
|
||
|
|
np.save(directory / "valid" / f"{session.name}.npy", mask)
|
||
|
|
targets[session.name] = np.asarray(session.targets, dtype=np.float32)
|
||
|
|
|
||
|
|
dense = {
|
||
|
|
name: evaluate_predictions(pred, valid, targets)
|
||
|
|
for name, pred in predictions.items()
|
||
|
|
}
|
||
|
|
endpoints = {
|
||
|
|
name: raw_endpoint_metrics(pred, valid, targets)
|
||
|
|
for name, pred in predictions.items()
|
||
|
|
}
|
||
|
|
for name, directory in directories.items():
|
||
|
|
metrics = {
|
||
|
|
**dense[name],
|
||
|
|
"format": "egospeak-omni-stage-two-eval-v1",
|
||
|
|
"classes": list(CLASS_NAMES),
|
||
|
|
"sessions": eval_names,
|
||
|
|
"train_sessions": train_names,
|
||
|
|
"checkpoint": str(checkpoint)
|
||
|
|
if name == "stage_two"
|
||
|
|
else provenance["stage_one_checkpoint"],
|
||
|
|
"method": name,
|
||
|
|
"seed": args.seed,
|
||
|
|
"prediction_layout": "[target_time, class, future_horizon]",
|
||
|
|
"raw_prediction_layout": "[source_time, future_horizon, class]",
|
||
|
|
"raw_endpoint_anticipation": endpoints[name],
|
||
|
|
}
|
||
|
|
if name == "stage_two":
|
||
|
|
metrics.update(
|
||
|
|
provenance=provenance,
|
||
|
|
budgets=budgets,
|
||
|
|
base_trainable=args.base_trainable,
|
||
|
|
training_config=training_config,
|
||
|
|
omni_sources={item.name: item.omni_metadata for item in sessions},
|
||
|
|
stage_one_reference=dense["stage_one"],
|
||
|
|
frozen_egospeak_baseline=dense["egospeak"],
|
||
|
|
semantic_evaluation="not run: provisional annotation audit pending",
|
||
|
|
)
|
||
|
|
(directory / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
|
||
|
|
print_metrics(name, dense[name])
|
||
|
|
csv_path = output / "raw_endpoint_metrics.csv"
|
||
|
|
with csv_path.open("w", newline="") as handle:
|
||
|
|
fields = [
|
||
|
|
"run",
|
||
|
|
"nominal_seconds",
|
||
|
|
"onset_endpoint_ap",
|
||
|
|
"offset_endpoint_ap",
|
||
|
|
"onset_positive",
|
||
|
|
"onset_negative",
|
||
|
|
"offset_positive",
|
||
|
|
"offset_negative",
|
||
|
|
]
|
||
|
|
writer = csv.DictWriter(handle, fieldnames=fields)
|
||
|
|
writer.writeheader()
|
||
|
|
for name, metrics in endpoints.items():
|
||
|
|
for item in metrics["horizons"]:
|
||
|
|
writer.writerow(
|
||
|
|
{
|
||
|
|
"run": name,
|
||
|
|
"nominal_seconds": item["seconds"],
|
||
|
|
"onset_endpoint_ap": item["onset"]["average_precision"],
|
||
|
|
"offset_endpoint_ap": item["offset"]["average_precision"],
|
||
|
|
"onset_positive": item["onset"]["num_positive"],
|
||
|
|
"onset_negative": item["onset"]["num_negative"],
|
||
|
|
"offset_positive": item["offset"]["num_positive"],
|
||
|
|
"offset_negative": item["offset"]["num_negative"],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
print("\nRaw endpoint ten-horizon means (percent; threshold-free):")
|
||
|
|
for name, metrics in endpoints.items():
|
||
|
|
print(
|
||
|
|
name,
|
||
|
|
{
|
||
|
|
k: None if metrics[k] is None else round(metrics[k] * 100, 3)
|
||
|
|
for k in ("mean_onset_ap", "mean_offset_ap")
|
||
|
|
},
|
||
|
|
)
|
||
|
|
print(f"Saved {checkpoint}; raw endpoint CSV (fractions)={csv_path}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|