"""Phase 9 — entraînement du modèle vidéo (images + audio synchronisés). Une vidéo = ``n_frames`` images + une piste audio synthétique. Le modèle encode **les frames comme une séquence** (ViT partagé + transformer spatio-temporel) et l'audio en parallèle, puis répond : * « Décris la vidéo : » -> « un cercle rouge va vers la droite » (frames) ; * « Que se passe-t-il ? » -> « un cercle rouge va vers la droite | voyelle a grave » (frames + audio synchronisés). Usage : python -m src.train.train_video --out-dir checkpoints/video_toy --epochs 400 """ from __future__ import annotations import argparse import json import time from dataclasses import asdict, dataclass from pathlib import Path import numpy as np import torch from src.data.audio_features import LogMelSpectrogram from src.data.toy_audio import SR from src.data.toy_audio import generate_dataset as gen_audio from src.data.toy_audio import to_arrays as audio_arrays from src.data.toy_video import generate_dataset as gen_video from src.data.toy_video import to_arrays as video_arrays from src.model.audio_encoder import AudioConfig, AudioEncoder from src.model.multimodal import VideoEncoder from src.model.projectors import MLPProjector from src.model.text_model import CortexTextLM, TextConfig from src.model.vit import ViTConfig from src.tokenizer import load_tokenizer from src.tokenizer.special_tokens import BOS_ID, EOS_ID, PAD_ID, SPECIAL_TOKEN_IDS from src.train.train_text import pick_device VID_S, VID_P, VID_E = (SPECIAL_TOKEN_IDS["<|video_start|>"], SPECIAL_TOKEN_IDS["<|video_pad|>"], SPECIAL_TOKEN_IDS["<|video_end|>"]) AUD_S, AUD_P, AUD_E = (SPECIAL_TOKEN_IDS["<|audio_start|>"], SPECIAL_TOKEN_IDS["<|audio_pad|>"], SPECIAL_TOKEN_IDS["<|audio_end|>"]) IGNORE_INDEX = -100 TASKS = ("video", "video_audio") class CortexVideoModel(torch.nn.Module): """Encodeur vidéo (ViT + spatio-temporel) + audio + décodeur texte.""" def __init__(self, text_cfg: TextConfig, vit_cfg: ViTConfig, audio_cfg: AudioConfig, n_frames: int): super().__init__() self.text = CortexTextLM(text_cfg) self.video = VideoEncoder(vit_cfg, n_frames, temporal_layers=2) self.audio = AudioEncoder(audio_cfg) self.video_projector = MLPProjector(vit_cfg.n_embd, text_cfg.n_embd) self.audio_projector = MLPProjector(audio_cfg.n_embd, text_cfg.n_embd) self.n_frames = n_frames @property def n_video_tokens(self) -> int: return self.n_frames * self.video.n_patches @property def n_audio_tokens(self) -> int: return self.audio.n_audio_tokens def video_features(self, frames): return self.video_projector(self.video(frames)) def audio_features(self, mels): return self.audio_projector(self.audio(mels)) def _assemble(self, input_ids, vmask, amask, frames, mels): emb = self.text.token_embedding(input_ids).clone() if frames is not None and bool(vmask.any()): rows = vmask.any(dim=1) feat = self.video_features(frames[rows]) emb[vmask] = feat.reshape(-1, feat.size(-1)) if mels is not None and bool(amask.any()): rows = amask.any(dim=1) feat = self.audio_features(mels[rows]) emb[amask] = feat.reshape(-1, feat.size(-1)) return emb def forward(self, input_ids, targets, vmask, amask, frames=None, mels=None): emb = self._assemble(input_ids, vmask, amask, frames, mels) return self.text(inputs_embeds=emb, targets=targets) @torch.no_grad() def generate(self, input_ids, vmask, amask, frames=None, mels=None, max_new_tokens=20, greedy=True, eos_id=None): self.eval() ids, vm, am = input_ids, vmask, amask for _ in range(max_new_tokens): emb = self._assemble(ids, vm, am, frames, mels) logits, _ = self.text(inputs_embeds=emb) logits = logits[:, -1, :] nxt = (torch.argmax(logits, dim=-1, keepdim=True) if greedy else torch.multinomial(torch.softmax(logits, dim=-1), 1)) ids = torch.cat([ids, nxt], dim=1) vm = torch.cat([vm, torch.zeros_like(nxt, dtype=torch.bool)], dim=1) am = torch.cat([am, torch.zeros_like(nxt, dtype=torch.bool)], dim=1) if eos_id is not None and (nxt == eos_id).all(): break return ids def make_batch(tokenizer, items, n_vid, n_aud, device): rows, vms, ams, labs = [], [], [], [] for task, item in items: prompt = "Décris la vidéo :" if task == "video" else "Que se passe-t-il ?" answer = (item[0].caption if task == "video" else f"{item[0].caption} | {item[1].caption}") ids = [BOS_ID, *tokenizer.encode(prompt)] vm = [False] * len(ids) am = [False] * len(ids) ids += [VID_S] + [VID_P] * n_vid + [VID_E] vm += [False] + [True] * n_vid + [False] am += [False] * (n_vid + 2) if task == "video_audio": ids += [AUD_S] + [AUD_P] * n_aud + [AUD_E] am += [False] + [True] * n_aud + [False] vm += [False] * (n_aud + 2) ans = [*tokenizer.encode(answer), EOS_ID] start = len(ids) ids += ans vm += [False] * len(ans) am += [False] * len(ans) lab = [IGNORE_INDEX] * len(ids) for t in range(start - 1, len(ids) - 1): lab[t] = ids[t + 1] rows.append(ids) vms.append(vm) ams.append(am) labs.append(lab) maxlen = max(len(r) for r in rows) for r, vm, am, lab in zip(rows, vms, ams, labs, strict=True): pad = maxlen - len(r) r += [PAD_ID] * pad vm += [False] * pad am += [False] * pad lab += [IGNORE_INDEX] * pad return (torch.tensor(rows, dtype=torch.long, device=device), torch.tensor(vms, dtype=torch.bool, device=device), torch.tensor(ams, dtype=torch.bool, device=device), torch.tensor(labs, dtype=torch.long, device=device)) @dataclass class VideoTrainConfig: out_dir: str = "checkpoints/video_toy" n_train: int = 1200 n_val: int = 120 grid: int = 32 n_frames: int = 4 epochs: int = 400 batch_size: int = 8 lr: float = 2e-3 weight_decay: float = 0.1 grad_clip: float = 1.0 text_layers: int = 2 text_embd: int = 96 text_heads: int = 4 vision_layers: int = 1 vision_embd: int = 96 patch_size: int = 8 audio_layers: int = 1 audio_embd: int = 96 n_mels: int = 64 n_frames_audio: int = 96 patch_mels: int = 16 patch_frames: int = 16 seed: int = 1337 device: str = "auto" def build_model(cfg: VideoTrainConfig, tokenizer): return CortexVideoModel( TextConfig(vocab_size=tokenizer.vocab_size, n_layer=cfg.text_layers, n_head=cfg.text_heads, n_embd=cfg.text_embd, ctx_len=300, tie_embeddings=True), ViTConfig(image_size=cfg.grid, patch_size=cfg.patch_size, n_layer=cfg.vision_layers, n_head=4, n_embd=cfg.vision_embd), AudioConfig(n_mels=cfg.n_mels, n_frames=cfg.n_frames_audio, patch_mels=cfg.patch_mels, patch_frames=cfg.patch_frames, n_layer=cfg.audio_layers, n_head=4, n_embd=cfg.audio_embd), cfg.n_frames, ) def _mels(items, cfg): tf = LogMelSpectrogram(sr=SR, n_fft=256, hop_length=64, n_mels=cfg.n_mels) with torch.no_grad(): m = tf(torch.from_numpy(audio_arrays(items))).squeeze(1) t = m.size(-1) m = m[:, :, :cfg.n_frames_audio] if t >= cfg.n_frames_audio else torch.nn.functional.pad( m, (0, cfg.n_frames_audio - t)) return m.numpy().astype(np.float32) def train(cfg: VideoTrainConfig) -> dict: torch.manual_seed(cfg.seed) device = pick_device(cfg.device) tokenizer = load_tokenizer() out_dir = Path(cfg.out_dir) out_dir.mkdir(parents=True, exist_ok=True) vtr = gen_video(cfg.n_train, grid=cfg.grid, n_frames=cfg.n_frames, seed=cfg.seed) vva = gen_video(cfg.n_val, grid=cfg.grid, n_frames=cfg.n_frames, seed=cfg.seed + 7) atr = gen_audio(cfg.n_train, seed=cfg.seed) ava = gen_audio(cfg.n_val, seed=cfg.seed + 7) Xtr, Xva = video_arrays(vtr), video_arrays(vva) Mtr, Mva = _mels(atr, cfg), _mels(ava, cfg) model = build_model(cfg, tokenizer).to(device) n_vid, n_aud = model.n_video_tokens, model.n_audio_tokens decay, no_decay = [], [] for _n, p in model.named_parameters(): if p.requires_grad: (decay if p.dim() >= 2 else no_decay).append(p) opt = torch.optim.AdamW( [{"params": decay, "weight_decay": cfg.weight_decay}, {"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr) rng = np.random.default_rng(cfg.seed) t0 = time.time() for epoch in range(cfg.epochs): idx = rng.integers(0, len(vtr), cfg.batch_size) tasks = rng.integers(0, 2, cfg.batch_size) items = [(TASKS[tasks[k]], (vtr[idx[k]], atr[idx[k]])) for k in range(len(idx))] ids, vm, am, lab = make_batch(tokenizer, items, n_vid, n_aud, device) frames = torch.from_numpy(Xtr[idx]).to(device) mels = torch.from_numpy(Mtr[idx]).unsqueeze(1).to(device) _, loss = model(ids, lab, vm, am, frames=frames, mels=mels) opt.zero_grad() loss.backward() if cfg.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip) opt.step() if (epoch + 1) % 50 == 0 or epoch == 0: print(f"epoch {epoch:4d} | loss {loss.item():.4f} | {time.time() - t0:.0f}s", flush=True) torch.save({"model": model.state_dict(), "train_config": asdict(cfg), "n_video_tokens": n_vid, "n_audio_tokens": n_aud}, out_dir / "best.pt") metrics = evaluate(model, tokenizer, Xva, Mva, vva, ava, n_vid, n_aud, device) summary = {"final_loss": round(loss.item(), 4), "epochs": cfg.epochs, "elapsed_s": round(time.time() - t0, 1), **metrics} (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") print("Résumé :", json.dumps({k: v for k, v in summary.items() if k != "examples"}, ensure_ascii=False), flush=True) return summary @torch.no_grad() def evaluate(model, tokenizer, Xva, Mva, vva, ava, n_vid, n_aud, device, n_eval=90): model.eval() n = min(n_eval, len(vva)) def tf(batch_items): ids, vm, am, lab = make_batch(tokenizer, batch_items, n_vid, n_aud, device) frames = torch.from_numpy(Xva[:n]).to(device) mels = torch.from_numpy(Mva[:n]).unsqueeze(1).to(device) logits, _ = model(ids, lab, vm, am, frames=frames, mels=mels) sup = lab != IGNORE_INDEX return round((logits.argmax(-1)[sup] == lab[sup]).float().mean().item(), 4) vid_items = [("video", (vva[i], ava[i])) for i in range(n)] both_items = [("video_audio", (vva[i], ava[i])) for i in range(n)] out = {"grounding_video": tf(vid_items), "grounding_video_audio": tf(both_items)} examples = [] for i in range(min(6, n)): ids = [BOS_ID, *tokenizer.encode("Que se passe-t-il ?")] vm = [False] * len(ids) am = [False] * len(ids) ids += [VID_S] + [VID_P] * n_vid + [VID_E] vm += [False] + [True] * n_vid + [False] am += [False] * (n_vid + 2) ids += [AUD_S] + [AUD_P] * n_aud + [AUD_E] am += [False] + [True] * n_aud + [False] vm += [False] * (n_aud + 2) out_ids = model.generate( torch.tensor([ids], device=device), torch.tensor([vm], dtype=torch.bool, device=device), torch.tensor([am], dtype=torch.bool, device=device), frames=torch.from_numpy(Xva[i:i + 1]).to(device), mels=torch.from_numpy(Mva[i:i + 1]).unsqueeze(1).to(device), max_new_tokens=24, greedy=True, eos_id=EOS_ID) text = tokenizer.decode(out_ids[0].tolist()[len(ids):], skip_special_tokens=True).strip() examples.append({"target": f"{vva[i].caption} | {ava[i].caption}", "pred": text}) correct = sum(1 for e in examples if e["pred"] == e["target"]) out["video_audio_greedy_acc"] = round(correct / len(examples), 4) out["examples"] = examples return out def main() -> None: parser = argparse.ArgumentParser(description="Entraîner le modèle vidéo (Phase 9).") for field, value in asdict(VideoTrainConfig()).items(): parser.add_argument(f"--{field.replace('_', '-')}", type=type(value), default=value) ns = parser.parse_args() train(VideoTrainConfig(**{k.replace("-", "_"): v for k, v in vars(ns).items()})) if __name__ == "__main__": main()