Add selective EgoSpeak Omni stage-two finetuning
This commit is contained in:
Executable
+520
@@ -0,0 +1,520 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user