Add transition-aware R1 and R2 objectives
This commit is contained in:
Xin Wang
+142
-4
@@ -28,6 +28,7 @@ from easycom_head import (
|
||||
normalize_sample_name,
|
||||
output_classes_for_loss,
|
||||
output_space_for_loss,
|
||||
transition_targets_from_activity,
|
||||
)
|
||||
|
||||
|
||||
@@ -112,8 +113,9 @@ def _predict_states(
|
||||
description: str,
|
||||
show_progress: bool,
|
||||
loss_name: str,
|
||||
) -> np.ndarray:
|
||||
) -> tuple[np.ndarray, np.ndarray | None]:
|
||||
chunks: list[torch.Tensor] = []
|
||||
transition_chunks: list[torch.Tensor] = []
|
||||
starts = range(0, len(states), batch_size)
|
||||
progress = tqdm(
|
||||
starts,
|
||||
@@ -128,10 +130,20 @@ def _predict_states(
|
||||
dtype=torch.float32,
|
||||
non_blocking=True,
|
||||
)
|
||||
if head.transition_auxiliary:
|
||||
activity_logits, transition_logits = head.forward_with_transition(batch)
|
||||
transition_chunks.append(transition_logits.sigmoid().cpu())
|
||||
else:
|
||||
activity_logits = head(batch)
|
||||
chunks.append(
|
||||
marginal_probabilities_from_logits(head(batch), loss_name).cpu()
|
||||
marginal_probabilities_from_logits(activity_logits, loss_name).cpu()
|
||||
)
|
||||
return torch.cat(chunks, dim=0).numpy()
|
||||
transition_probabilities = (
|
||||
torch.cat(transition_chunks, dim=0).numpy()
|
||||
if transition_chunks
|
||||
else None
|
||||
)
|
||||
return torch.cat(chunks, dim=0).numpy(), transition_probabilities
|
||||
|
||||
|
||||
def _align_to_egospeak_timeline(
|
||||
@@ -156,6 +168,82 @@ def _align_to_egospeak_timeline(
|
||||
return predictions, valid
|
||||
|
||||
|
||||
def _align_transitions_to_timeline(
|
||||
probabilities: np.ndarray,
|
||||
label_indices: torch.LongTensor,
|
||||
num_target_steps: int,
|
||||
) -> np.ndarray:
|
||||
"""Align [current_state, future_horizon] transition probabilities."""
|
||||
|
||||
predictions = np.zeros(
|
||||
(num_target_steps, NUM_HORIZONS),
|
||||
dtype=np.float32,
|
||||
)
|
||||
current = label_indices.numpy()
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
target_indices = current + horizon + 1
|
||||
usable = (target_indices >= 0) & (target_indices < num_target_steps)
|
||||
predictions[target_indices[usable], horizon] = probabilities[usable, horizon]
|
||||
return predictions
|
||||
|
||||
|
||||
def _evaluate_transitions(
|
||||
session_predictions: dict[str, np.ndarray],
|
||||
session_valid: dict[str, np.ndarray],
|
||||
session_targets: dict[str, np.ndarray],
|
||||
) -> dict[str, Any]:
|
||||
"""Evaluate auxiliary transition probabilities on causal-valid rows."""
|
||||
|
||||
from sklearn.metrics import average_precision_score
|
||||
|
||||
horizons = []
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
scores = np.concatenate(
|
||||
[
|
||||
session_predictions[name][session_valid[name][:, horizon], horizon]
|
||||
for name in session_predictions
|
||||
]
|
||||
)
|
||||
targets = np.concatenate(
|
||||
[
|
||||
session_targets[name][session_valid[name][:, horizon]]
|
||||
for name in session_predictions
|
||||
]
|
||||
)
|
||||
predicted = scores >= 0.5
|
||||
positive = targets == 1
|
||||
true_positive = int(np.logical_and(predicted, positive).sum())
|
||||
false_positive = int(np.logical_and(predicted, ~positive).sum())
|
||||
false_negative = int(np.logical_and(~predicted, positive).sum())
|
||||
precision = true_positive / max(1, true_positive + false_positive)
|
||||
recall = true_positive / max(1, true_positive + false_negative)
|
||||
f1 = 2.0 * precision * recall / max(1e-12, precision + recall)
|
||||
horizons.append(
|
||||
{
|
||||
"horizon": horizon + 1,
|
||||
"seconds": (horizon + 1) * STEP_SECONDS,
|
||||
"num_steps": int(len(scores)),
|
||||
"positive_rate": float(positive.mean()),
|
||||
"average_precision": float(
|
||||
average_precision_score(targets, scores)
|
||||
),
|
||||
"precision_at_0.5": precision,
|
||||
"recall_at_0.5": recall,
|
||||
"f1_at_0.5": f1,
|
||||
}
|
||||
)
|
||||
return {
|
||||
"definition": (
|
||||
"A positive at target row u means the three-label activity state "
|
||||
"changed between rows u-1 and u."
|
||||
),
|
||||
"horizons": horizons,
|
||||
"mean_average_precision": float(
|
||||
np.mean([item["average_precision"] for item in horizons])
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _evaluate(
|
||||
session_predictions: dict[str, np.ndarray],
|
||||
session_valid: dict[str, np.ndarray],
|
||||
@@ -278,6 +366,11 @@ def main() -> None:
|
||||
session_predictions: dict[str, np.ndarray] = {}
|
||||
session_valid: dict[str, np.ndarray] = {}
|
||||
session_targets: dict[str, np.ndarray] = {}
|
||||
session_transition_predictions: dict[str, np.ndarray] = {}
|
||||
session_transition_targets: dict[str, np.ndarray] = {}
|
||||
transition_prediction_dir = args.output_dir / "transition_pred"
|
||||
if head.transition_auxiliary:
|
||||
transition_prediction_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"device: {device}")
|
||||
print(
|
||||
@@ -302,7 +395,7 @@ def main() -> None:
|
||||
f"{name} has {data.num_input_layers} feature layers, checkpoint "
|
||||
f"expects {head_config['num_input_layers']}"
|
||||
)
|
||||
probabilities = _predict_states(
|
||||
probabilities, transition_probabilities = _predict_states(
|
||||
head,
|
||||
data.states,
|
||||
device,
|
||||
@@ -322,12 +415,37 @@ def main() -> None:
|
||||
session_targets[name] = target_array
|
||||
np.save(prediction_dir / f"{name}.npy", predictions)
|
||||
np.save(valid_dir / f"{name}.npy", valid)
|
||||
if transition_probabilities is not None:
|
||||
transition_predictions = _align_transitions_to_timeline(
|
||||
transition_probabilities,
|
||||
data.label_indices,
|
||||
len(data.targets),
|
||||
)
|
||||
transition_targets = np.zeros(len(data.targets), dtype=np.float32)
|
||||
transition_targets[1:] = (
|
||||
transition_targets_from_activity(data.targets).numpy()
|
||||
)
|
||||
session_transition_predictions[name] = transition_predictions
|
||||
session_transition_targets[name] = transition_targets
|
||||
np.save(
|
||||
transition_prediction_dir / f"{name}.npy",
|
||||
transition_predictions,
|
||||
)
|
||||
print(
|
||||
f"{name}: saved pred {predictions.shape}, "
|
||||
f"causal-valid counts={valid.sum(axis=0).tolist()}"
|
||||
)
|
||||
|
||||
metrics = _evaluate(session_predictions, session_valid, session_targets)
|
||||
transition_metrics = (
|
||||
_evaluate_transitions(
|
||||
session_transition_predictions,
|
||||
session_valid,
|
||||
session_transition_targets,
|
||||
)
|
||||
if session_transition_predictions
|
||||
else None
|
||||
)
|
||||
metrics.update(
|
||||
{
|
||||
"sessions": sessions,
|
||||
@@ -346,6 +464,12 @@ def main() -> None:
|
||||
"checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
|
||||
"checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
|
||||
"layer_fusion": head.layer_fusion,
|
||||
"transition_auxiliary": head.transition_auxiliary,
|
||||
"transition_objective": checkpoint_payload.get(
|
||||
"training_config",
|
||||
{},
|
||||
).get("transition_objective", "none"),
|
||||
"transition_metrics": transition_metrics,
|
||||
"layer_mixture_weights_oldest_to_newest": (
|
||||
head.layer_mixture_weights().detach().float().cpu().tolist()
|
||||
if head.layer_mixture_weights() is not None
|
||||
@@ -356,6 +480,20 @@ def main() -> None:
|
||||
metrics_path = args.output_dir / "metrics.json"
|
||||
metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
|
||||
_print_metrics(metrics)
|
||||
if transition_metrics is not None:
|
||||
print("\ntransition_auxiliary:")
|
||||
for result in transition_metrics["horizons"]:
|
||||
print(
|
||||
f"{result['seconds']:.1f}s: "
|
||||
f"AP={result['average_precision'] * 100:.5f}, "
|
||||
f"P@0.5={result['precision_at_0.5'] * 100:.5f}, "
|
||||
f"R@0.5={result['recall_at_0.5'] * 100:.5f}, "
|
||||
f"F1@0.5={result['f1_at_0.5'] * 100:.5f}"
|
||||
)
|
||||
print(
|
||||
"10-horizon transition mean AP: "
|
||||
f"{transition_metrics['mean_average_precision'] * 100:.5f}"
|
||||
)
|
||||
if metrics["layer_mixture_weights_oldest_to_newest"] is not None:
|
||||
print(
|
||||
"\nLayer mixture weights (oldest->newest): "
|
||||
|
||||
Reference in New Issue
Block a user