| """``nida train``: config + latent dataset directory -> a (resumable) training run. |
| |
| Workspace layout under ``--out``:: |
| |
| checkpoints/ckpt-<step>.pt atomic, newest ``keep_last`` kept |
| metrics.jsonl train/eval records (replayed steps are removed on resume) |
| run_info.json config, seed, git commit, data fingerprint, load reports, |
| ``data_dirs`` (all dataset directories, in load order), |
| ``dataset_meta`` (the single directory's ``meta.json``, null |
| for several directories) and ``dataset_metas`` (one |
| ``meta.json`` or null per directory) |
| (``resume_info-<step>.json`` when a run is resumed) |
| data_rejections.jsonl every item rejected at load time, with the reason |
| latent_stats.json per-channel mean/std the model was trained on (inference |
| needs it to de-normalise generated latents) |
| config.yaml verbatim copy of the training config file |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import logging |
| import subprocess |
| from collections.abc import Sequence |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import torch |
|
|
| from nida.train.checkpoint import list_checkpoints |
| from nida.train.config import TrainConfig, load_train_config |
| from nida.train.data import ( |
| LANGUAGES, |
| DataFormatError, |
| LatentDataset, |
| LoadReport, |
| load_latent_datasets, |
| ) |
| from nida.train.loop import CHECKPOINT_DIR, ResumeError, Trainer, TrainResult |
| from nida.utils.io import atomic_write_json, atomic_write_text |
|
|
| __all__ = ["TrainOutcome", "check_languages", "render_report", "resolve_git_commit", "run_training"] |
|
|
| log = logging.getLogger(__name__) |
|
|
|
|
| @dataclass |
| class TrainOutcome: |
| """What ``run_training`` returns: the raw result plus the JSON/markdown summaries.""" |
|
|
| result: TrainResult |
| metrics: dict[str, object] |
| report_md: str |
|
|
|
|
| def resolve_git_commit(explicit: str | None, repo_dir: Path) -> str | None: |
| """``explicit`` if given, else ``git rev-parse HEAD`` in ``repo_dir`` (None if unavailable).""" |
| if explicit is not None: |
| return explicit |
| try: |
| out = subprocess.run( |
| ["git", "rev-parse", "HEAD"], |
| cwd=repo_dir, |
| capture_output=True, |
| text=True, |
| check=True, |
| timeout=10, |
| ) |
| except (OSError, subprocess.SubprocessError): |
| return None |
| return out.stdout.strip() or None |
|
|
|
|
| def _rejection_lines(reports: list[LoadReport]) -> str: |
| lines = [ |
| json.dumps( |
| {"split": r.split, "segment_id": x.segment_id, "reason": x.reason, "detail": x.detail}, |
| ensure_ascii=False, |
| sort_keys=True, |
| ) |
| for r in reports |
| for x in r.rejections |
| ] |
| return "".join(f"{line}\n" for line in lines) |
|
|
|
|
| def _read_dataset_meta(data_dir: Path) -> object: |
| """The directory's ``meta.json`` content, or None when it has none.""" |
| meta_path = data_dir / "meta.json" |
| if not meta_path.is_file(): |
| return None |
| return json.loads(meta_path.read_text(encoding="utf-8")) |
|
|
|
|
| def _pick_device(name: str | None) -> torch.device: |
| if name is None or name == "auto": |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| return torch.device(name) |
|
|
|
|
| def check_languages(n_languages: int, datasets: Sequence[LatentDataset]) -> None: |
| """Refuse data whose languages the model cannot represent (ADR-0007), before training. |
| |
| ``n_languages == 0`` (no language embedding) accepts only Turkish (id 0); otherwise every |
| segment's language id must be below ``n_languages``. |
| |
| Raises: |
| DataFormatError: naming the first offending language and how many segments use it. |
| """ |
| allowed = max(n_languages, 1) |
| counts: dict[str, int] = {} |
| for data in datasets: |
| for seg in data.segments: |
| if LANGUAGES.index(seg.language) >= allowed: |
| counts[seg.language] = counts.get(seg.language, 0) + 1 |
| if counts: |
| raise DataFormatError( |
| f"model n_languages={n_languages} cannot represent languages {counts} " |
| f"(order {LANGUAGES}); set model.overrides.n_languages" |
| ) |
|
|
|
|
| def run_training( |
| config_path: Path, |
| data_dirs: Path | Sequence[Path], |
| out_dir: Path, |
| *, |
| resume: bool, |
| device: str | None = None, |
| git_commit: str | None = None, |
| repo_dir: Path | None = None, |
| handle_sigterm: bool = True, |
| init_from: Path | None = None, |
| ) -> TrainOutcome: |
| """Load config and data, train, and summarise. |
| |
| ``data_dirs`` is one latent dataset directory or several, trained on as their union |
| (:func:`nida.train.data.load_latent_datasets`); latent stats come from the merged |
| train split. |
| |
| ``resume=True`` continues from the newest checkpoint in ``out_dir`` (starting fresh, |
| with a warning, when there is none -- so a restart loop can always pass ``--resume``). |
| ``resume=False`` refuses to run when checkpoints already exist. |
| |
| ``init_from`` (a checkpoint of another run with the same model config, or the parent of |
| a ``dur_cond`` "none" -> "emotion" step, see :meth:`Trainer.init_from`) starts a fresh |
| run -- a fine-tune -- from its weights and latent stats (:meth:`Trainer.init_from`). It |
| applies only when no checkpoint of this run was resumed, so ``--resume --init-from`` |
| restarts safely: the first launch inherits the weights, later launches continue the |
| fine-tune's own checkpoints. |
| |
| Raises: |
| ConfigError, DataFormatError, BudgetError, ResumeError, CheckpointError: see the |
| respective modules; nothing is swallowed. |
| """ |
| cfg = load_train_config(config_path) |
| dirs = [data_dirs] if isinstance(data_dirs, Path) else list(data_dirs) |
| dirs_label = ", ".join(str(d) for d in dirs) |
| |
| timings = cfg.model_config().text_path == "aligned" |
| train_data, train_report = load_latent_datasets(dirs, "train", timings=timings) |
| try: |
| val_data, val_report = load_latent_datasets(dirs, "val", timings=timings) |
| except Exception: |
| train_data.close() |
| raise |
| try: |
| if len(train_data) == 0: |
| raise DataFormatError( |
| f"no usable training items in {dirs_label} " |
| f"(rejections: {train_report.rejection_counts})" |
| ) |
| check_languages(cfg.model_config().n_languages, [train_data, val_data]) |
| if len(val_data) == 0 and cfg.eval_every is not None: |
| log.warning( |
| "no usable validation items in %s: evaluation (eval_every=%s) is skipped", |
| dirs_label, |
| cfg.eval_every, |
| ) |
| checkpoint_dir = out_dir / CHECKPOINT_DIR |
| existing = list_checkpoints(checkpoint_dir) |
| if existing and not resume: |
| raise ResumeError( |
| f"{checkpoint_dir} already holds checkpoints (latest step {existing[-1][0]}); " |
| "pass --resume to continue or choose a fresh --out" |
| ) |
| commit = resolve_git_commit(git_commit, repo_dir or Path.cwd()) |
| trainer = Trainer( |
| cfg, |
| train_data, |
| out_dir, |
| val_data=val_data, |
| device=_pick_device(device), |
| git_commit=commit, |
| ) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| atomic_write_text( |
| out_dir / "data_rejections.jsonl", _rejection_lines([train_report, val_report]) |
| ) |
| atomic_write_text(out_dir / "config.yaml", config_path.read_text(encoding="utf-8")) |
| dataset_metas = [_read_dataset_meta(d) for d in dirs] |
| resumed_path = trainer.resume_latest() if resume else None |
| if resume and resumed_path is None: |
| log.warning("--resume given but %s has no checkpoint; starting fresh", checkpoint_dir) |
| if resumed_path is None and init_from is not None: |
| trainer.init_from(init_from) |
| elif init_from is not None and trainer.init_from_path != str(init_from.resolve()): |
| log.warning( |
| "--init-from %s ignored: resumed %s, whose run started from %s", |
| init_from, |
| resumed_path, |
| trainer.init_from_path, |
| ) |
| atomic_write_json(out_dir / "latent_stats.json", trainer.latent_stats.to_dict()) |
| info_name = ( |
| "run_info.json" if resumed_path is None else f"resume_info-{trainer.step:09d}.json" |
| ) |
| atomic_write_json( |
| out_dir / info_name, |
| { |
| "config": cfg.to_dict(), |
| "seed": cfg.seed, |
| "git_commit": commit, |
| "device": str(trainer.device), |
| "dataset_fingerprint": train_data.fingerprint(), |
| "resumed_from": None if resumed_path is None else str(resumed_path), |
| "init_from": trainer.init_from_path, |
| "init_from_model_diff": trainer.init_from_model_diff, |
| "latent_stats_inherited": trainer.latent_stats_inherited, |
| "data": {"train": train_report.to_dict(), "val": val_report.to_dict()}, |
| "data_dirs": [str(d) for d in dirs], |
| "dataset_meta": dataset_metas[0] if len(dirs) == 1 else None, |
| "dataset_metas": dataset_metas, |
| "prompt_pairing": { |
| **cfg.prompt_pairing.to_dict(), |
| **trainer.speaker_dirs.to_dict(), |
| }, |
| }, |
| ) |
| result = trainer.run(handle_sigterm=handle_sigterm) |
| num_params = sum(p.numel() for p in trainer.model.parameters()) |
| metrics = _metrics(cfg, result, train_report, val_report, trainer, commit, num_params) |
| finally: |
| _close(train_data, val_data) |
| return TrainOutcome(result, metrics, render_report(metrics)) |
|
|
|
|
| def _close(*datasets: LatentDataset) -> None: |
| for dataset in datasets: |
| dataset.close() |
|
|
|
|
| def _metrics( |
| cfg: TrainConfig, |
| result: TrainResult, |
| train_report: LoadReport, |
| val_report: LoadReport, |
| trainer: Trainer, |
| commit: str | None, |
| num_params: int, |
| ) -> dict[str, object]: |
| return { |
| "completed": result.completed, |
| "stopped_by_signal": result.stopped_by_signal, |
| "steps_done": result.steps_done, |
| "total_steps": cfg.steps, |
| "last_loss": result.last_loss, |
| "mean_frames_per_s": result.mean_frames_per_s, |
| "resumed_from_step": result.resumed_from_step, |
| "evals": result.evals, |
| "checkpoints": result.checkpoints, |
| "best_checkpoint": result.best_checkpoint, |
| "best_val_loss": trainer.best_val_loss, |
| "seed": cfg.seed, |
| "git_commit": commit, |
| "device": str(trainer.device), |
| "model": cfg.model.name, |
| "num_params": num_params, |
| "overfit_one_batch": cfg.overfit_one_batch, |
| |
| |
| "prompt_pairing": { |
| **cfg.prompt_pairing.to_dict(), |
| **trainer.speaker_dirs.to_dict(), |
| "totals": trainer.pairing_totals.to_dict(), |
| }, |
| "dataset_fingerprint": trainer.train_data.fingerprint(), |
| "out_dir": str(trainer.out_dir), |
| "data": {"train": train_report.to_dict(), "val": val_report.to_dict()}, |
| } |
|
|
|
|
| def render_report(metrics: dict[str, object]) -> str: |
| """Markdown run report from the ``metrics`` dict (a human-readable summary).""" |
|
|
| def show(value: object) -> str: |
| return "-" if value is None else str(value) |
|
|
| state = ( |
| "completed" |
| if metrics["completed"] |
| else ( |
| "INTERRUPTED by SIGTERM (checkpoint saved, resume with --resume)" |
| if metrics["stopped_by_signal"] |
| else "stopped early" |
| ) |
| ) |
| pairing = metrics["prompt_pairing"] |
| assert isinstance(pairing, dict) |
| totals = pairing["totals"] |
| pairing_line = ( |
| "off (in-clip prompt)" |
| if pairing["cross_p"] == 0 |
| else f"cross_p={pairing['cross_p']}, prompt {pairing['min_prompt_frames']}.." |
| f"{pairing['max_prompt_frames']} frames, speaker dirs {pairing['matched_dirs']}; " |
| f"{totals['paired']} of {totals['items']} items cross-segment, " |
| f"fallbacks {totals['fallbacks']}" |
| ) |
| lines = [ |
| "# train", |
| "", |
| f"- state: {state}", |
| f"- steps: {metrics['steps_done']} / {metrics['total_steps']}", |
| f"- resumed from step: {show(metrics['resumed_from_step'])}", |
| f"- last loss: {show(metrics['last_loss'])}", |
| f"- mean frames/s (this session): {show(metrics['mean_frames_per_s'])}", |
| f"- model: {metrics['model']} ({metrics['num_params']} parameters)", |
| f"- device: {metrics['device']}", |
| f"- seed: {metrics['seed']}", |
| f"- git commit: {show(metrics['git_commit'])}", |
| f"- overfit_one_batch: {metrics['overfit_one_batch']}", |
| f"- prompt pairing: {pairing_line}", |
| f"- dataset fingerprint: {metrics['dataset_fingerprint']}", |
| f"- workspace: {metrics['out_dir']}", |
| "", |
| "## Data", |
| "", |
| ] |
| data = metrics["data"] |
| assert isinstance(data, dict) |
| for split in ("train", "val"): |
| rep = data[split] |
| lines.append( |
| f"- {split}: {rep['num_accepted']} accepted of {rep['num_in_split']} " |
| f"({rep['accepted_frames']} frames), rejected {rep['num_rejected']} " |
| f"{rep['rejection_counts']}" |
| ) |
| evals = metrics["evals"] |
| assert isinstance(evals, list) |
| if evals: |
| lines += [ |
| "", |
| "## Validation", |
| "", |
| "| step | val_loss | val_loss_ema | val_loss_ema (no text) |", |
| "|---|---|---|---|", |
| ] |
| lines += [ |
| f"| {int(e['step'])} | {e['val_loss']:.5f} | {e['val_loss_ema']:.5f} | " |
| f"{e.get('val_loss_ema_notext', float('nan')):.5f} |" |
| for e in evals |
| ] |
| lines += [ |
| "", |
| f"- best val_loss_ema: {show(metrics['best_val_loss'])} " |
| f"({show(metrics['best_checkpoint'])})", |
| ] |
| checkpoints = metrics["checkpoints"] |
| assert isinstance(checkpoints, list) |
| if checkpoints: |
| lines += ["", "## Checkpoints written this session", ""] |
| lines += [f"- {c}" for c in checkpoints] |
| return "\n".join(lines) + "\n" |
|
|
