Add selective EgoSpeak Omni stage-two finetuning
This commit is contained in:
Xin Wang
@@ -0,0 +1,195 @@
|
||||
"""Stage-two selective MiniROADA adaptation from a trained Omni residual.
|
||||
|
||||
The released input features and Omni states remain frozen. This module only
|
||||
enables gradients in selected MiniROADA temporal/forecast modules and the
|
||||
existing residual. Stage-one checkpoints are loaded strictly, without resetting
|
||||
the learned residual or requiring the original checkpoint to be available.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
|
||||
from easycom_head import NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS
|
||||
from egospeak_omni_residual import (
|
||||
EgoSpeakOmniResidual,
|
||||
MiniROADABase,
|
||||
load_torch_payload,
|
||||
normalize_session_name,
|
||||
)
|
||||
|
||||
BaseTrainable = Literal["none", "heads", "gru-heads"]
|
||||
STAGE_ONE_FORMAT = "egospeak-omni-zero-residual-v1"
|
||||
STAGE_TWO_FORMAT = "egospeak-omni-selective-finetune-v1"
|
||||
|
||||
|
||||
def validate_training_sessions(values: list[str]) -> None:
|
||||
normalized = [normalize_session_name(value) for value in values]
|
||||
if not normalized or len(normalized) != len(set(normalized)):
|
||||
raise ValueError("Training sessions must be nonempty and unique")
|
||||
if not set(normalized).issubset(OFFICIAL_TRAIN_SESSIONS):
|
||||
raise ValueError("Training may only use EasyCom Sessions 4--12")
|
||||
|
||||
|
||||
def load_stage_one_checkpoint(
|
||||
path: Path,
|
||||
) -> tuple[EgoSpeakOmniResidual, dict[str, Any]]:
|
||||
path = path.expanduser().resolve()
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(path)
|
||||
payload = load_torch_payload(path)
|
||||
if not isinstance(payload, dict) or payload.get("format") != STAGE_ONE_FORMAT:
|
||||
raise ValueError(
|
||||
"--stage-one-checkpoint must be a stage-one residual checkpoint"
|
||||
)
|
||||
config = payload.get("model_config")
|
||||
if not isinstance(config, dict):
|
||||
raise TypeError("Stage-one checkpoint lacks model_config")
|
||||
required = {
|
||||
"omni_dim",
|
||||
"fusion_dim",
|
||||
"tail_embedding_dim",
|
||||
"num_tail_bins",
|
||||
"dropout",
|
||||
}
|
||||
if set(config) != required or config["num_tail_bins"] != NUM_HORIZONS:
|
||||
raise ValueError("Incompatible stage-one residual model_config")
|
||||
training = payload.get("training_config")
|
||||
if not isinstance(training, dict) or int(training.get("epochs", 0)) <= 0:
|
||||
raise ValueError("Stage-one checkpoint lacks a positive training epoch budget")
|
||||
validate_training_sessions(payload.get("train_sessions", []))
|
||||
base_config = payload.get("base_model_config", {})
|
||||
if not isinstance(base_config, dict):
|
||||
raise TypeError("Invalid base_model_config")
|
||||
# Released v1 checkpoints omit base_model_config and use default MiniROADA.
|
||||
model = EgoSpeakOmniResidual(MiniROADABase(**base_config), **config)
|
||||
model.load_state_dict(payload["model_state_dict"], strict=True)
|
||||
if any(not torch.isfinite(value).all() for value in model.state_dict().values()):
|
||||
raise ValueError("Stage-one checkpoint contains nonfinite parameters")
|
||||
model.requires_grad_(False)
|
||||
model.eval()
|
||||
return model, payload
|
||||
|
||||
|
||||
class EgoSpeakOmniFinetune(EgoSpeakOmniResidual):
|
||||
"""Continuation with explicit base policy and intact stage-one logits."""
|
||||
|
||||
def __init__(
|
||||
self, *args: Any, base_trainable: BaseTrainable = "gru-heads", **kwargs: Any
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
if base_trainable not in ("none", "heads", "gru-heads"):
|
||||
raise ValueError(f"Invalid base-trainable policy: {base_trainable}")
|
||||
self.base_trainable = base_trainable
|
||||
self.requires_grad_(True)
|
||||
self.base.requires_grad_(False)
|
||||
if base_trainable in ("heads", "gru-heads"):
|
||||
self.base.anticipation_layer.requires_grad_(True)
|
||||
self.base.f_classification.requires_grad_(True)
|
||||
if base_trainable == "gru-heads":
|
||||
self.base.gru.requires_grad_(True)
|
||||
self.train(False)
|
||||
|
||||
@classmethod
|
||||
def from_stage_one(
|
||||
cls,
|
||||
reference: EgoSpeakOmniResidual,
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
base_trainable: BaseTrainable = "gru-heads",
|
||||
) -> EgoSpeakOmniFinetune:
|
||||
model = cls(deepcopy(reference.base), **config, base_trainable=base_trainable)
|
||||
# Copy rather than share weights with the immutable first-stage reference.
|
||||
model.load_state_dict(reference.state_dict(), strict=True)
|
||||
return model
|
||||
|
||||
def train(self, mode: bool = True) -> EgoSpeakOmniFinetune:
|
||||
super().train(mode)
|
||||
# Keep frozen layer1's dropout off. Enabling GRU/head train mode does
|
||||
# not turn the original frozen RGB+audio projection back on.
|
||||
if self.base_trainable in ("heads", "gru-heads"):
|
||||
self.base.anticipation_layer.train(mode)
|
||||
self.base.f_classification.train(mode)
|
||||
if self.base_trainable == "gru-heads":
|
||||
self.base.gru.train(mode)
|
||||
return self
|
||||
|
||||
def optimizer_groups(
|
||||
self, *, fusion_lr: float, base_lr: float
|
||||
) -> list[dict[str, Any]]:
|
||||
if fusion_lr <= 0 or base_lr <= 0:
|
||||
raise ValueError("Both learning rates must be positive")
|
||||
fusion = [
|
||||
p
|
||||
for name, p in self.named_parameters()
|
||||
if p.requires_grad and not name.startswith("base.")
|
||||
]
|
||||
base = [
|
||||
p
|
||||
for name, p in self.named_parameters()
|
||||
if p.requires_grad and name.startswith("base.")
|
||||
]
|
||||
groups = [{"name": "fusion", "params": fusion, "lr": fusion_lr}]
|
||||
if base:
|
||||
groups.append({"name": "miniroada", "params": base, "lr": base_lr})
|
||||
return groups
|
||||
|
||||
def forward_last(
|
||||
self,
|
||||
rgb: torch.Tensor,
|
||||
audio: torch.Tensor,
|
||||
omni: torch.Tensor,
|
||||
tail_bin: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
hidden, _ = self.base.encode(rgb, audio)
|
||||
ego_hidden = hidden[:, -1, :]
|
||||
base_logits = self.base.anticipation_logits(ego_hidden)
|
||||
# No no_grad/detach: gradients must reach the selected MiniROADA GRU.
|
||||
residual = self._residual_logits(ego_hidden, omni, tail_bin)
|
||||
return base_logits + residual, base_logits
|
||||
|
||||
def forward_sequence_chunk(
|
||||
self,
|
||||
rgb: torch.Tensor,
|
||||
audio: torch.Tensor,
|
||||
omni: torch.Tensor,
|
||||
tail_bin: torch.Tensor,
|
||||
recurrent_state: torch.Tensor | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
hidden, next_state = self.base.encode(rgb, audio, recurrent_state)
|
||||
base_logits = self.base.anticipation_logits(hidden)
|
||||
residual = self._residual_logits(hidden, omni, tail_bin)
|
||||
return base_logits + residual, base_logits, next_state
|
||||
|
||||
|
||||
def base_model_config(base: MiniROADABase) -> dict[str, Any]:
|
||||
return {
|
||||
"rgb_dim": base.rgb_dim,
|
||||
"audio_dim": base.audio_dim,
|
||||
"embedding_dim": base.embedding_dim,
|
||||
"hidden_dim": base.hidden_dim,
|
||||
"num_layers": base.num_layers,
|
||||
"num_horizons": base.num_horizons,
|
||||
"num_classes": base.num_classes,
|
||||
"dropout": base.layer1[-1].p,
|
||||
}
|
||||
|
||||
|
||||
def load_stage_two_checkpoint(
|
||||
path: Path,
|
||||
) -> tuple[EgoSpeakOmniFinetune, dict[str, Any]]:
|
||||
payload = load_torch_payload(path)
|
||||
if not isinstance(payload, dict) or payload.get("format") != STAGE_TWO_FORMAT:
|
||||
raise ValueError("Not a stage-two fine-tuning checkpoint")
|
||||
model = EgoSpeakOmniFinetune(
|
||||
MiniROADABase(**payload["base_model_config"]),
|
||||
**payload["model_config"],
|
||||
base_trainable=payload["base_trainable"],
|
||||
)
|
||||
model.load_state_dict(payload["model_state_dict"], strict=True)
|
||||
model.eval()
|
||||
return model, payload
|
||||
Reference in New Issue
Block a user