vmaf-train — tiny-AI training harness CLI¶

vmaf-train trains, exports, audits and registers the tiny ONNX models that vmaf and vmaf-tune consume. It is the Python entry point of the fork's tiny-AI training infrastructure (ai/src/vmaf_train/cli.py, registered in ai/pyproject.toml [project.scripts]). This page covers all 15 subcommands. For what the models do, see docs/ai/overview.md; for the training-corpus pipeline, see docs/ai/training.md.

Install¶

pip install -e ai
vmaf-train --help

Quick path¶

The pipeline runs in this order: extract features, fit (or tune), export, check ops, validate normalisation, evaluate, register. The numbered recipe below trains and registers a new fr_regressor model.

1. Pre-compute features for a dataset named in the manifest (nflx, konvid-1k, live-vqc, youtube-ugc or bvi-dvc):

   vmaf-train extract-features \
     --dataset nflx \
     --output .corpus/netflix/features.parquet
   

2. Train from a YAML config. --output is the run directory; the checkpoint is <output>/last.ckpt. Use tune instead for an Optuna sweep:

   vmaf-train fit \
     --config ai/configs/fr_tiny_v1.yaml \
     --cache .corpus/netflix/features.parquet \
     --output runs/fr_tiny_v1
   

3. Export to ONNX:

   vmaf-train export \
     --checkpoint runs/fr_tiny_v1/last.ckpt \
     --output runs/fr_tiny_v1/fr_tiny_v1.onnx \
     --model fr_regressor
   

4. Check the op allowlist:

   vmaf-train check-ops --model runs/fr_tiny_v1/fr_tiny_v1.onnx
   

5. Validate the normalisation against the corpus:

   vmaf-train validate-norm \
     --model runs/fr_tiny_v1/fr_tiny_v1.onnx \
     --features .corpus/netflix/features.parquet \
     --fail-on-warning
   

6. Evaluate on the test split:

   vmaf-train eval \
     --model runs/fr_tiny_v1/fr_tiny_v1.onnx \
     --features .corpus/netflix/features.parquet \
     --split test
   

7. Register the model (copy it under model/tiny/ first):

   vmaf-train register \
     --model model/tiny/fr_tiny_v1.onnx \
     --kind fr \
     --dataset nflx \
     --license BSD-2-Clause-Patent \
     --train-commit "$(git rev-parse HEAD)" \
     --train-config ai/configs/fr_tiny_v1.yaml
   

To shrink a trained model, quantise it:

vmaf-train quantize-int8 \
  --fp32 model/tiny/fr_regressor_v1.onnx \
  --output model/tiny/fr_regressor_v1.int8.onnx \
  --calibration .corpus/netflix/features.parquet \
  --rmse-gate 0.5 \
  --json quantize-report.json

Subcommand reference¶

Train¶

extract-features¶

Pre-compute libvmaf features over a dataset's (ref, dis) pairs and write them to a parquet cache.

manifest-scan¶

Scan a local corpus directory for YUV, Y4M, MP4, MKV and WebM files, pin each by SHA-256, and overwrite the in-tree manifest of the dataset.

fit¶

Train a model from a YAML config. Configs live in ai/configs/ (fr_tiny_v1.yaml, nr_mobilenet_v1.yaml, filter_residual_v1.yaml, learned_filter_v1_qat.yaml).

tune¶

Run an Optuna sweep over the model_args of a base YAML config.

Export and evaluate¶

export¶

Export a trained checkpoint to ONNX with a round-trip validation step.

eval¶

Evaluate an ONNX model on a deterministic split and report PLCC, SROCC and RMSE.

quantize-int8¶

Post-training static int8 quantisation (QDQ format) per ADR-0173. It reports drift against held-out samples and exits 2 if the RMSE breaks the gate.

Audit¶

All audit commands that take --json write the report to the given path (they do not switch stdout to JSON).

check-ops¶

Check one ONNX model against libvmaf's op allowlist (core/src/dnn/op_allowlist.c). Exits 2 if forbidden ops are found.

audit-compat¶

Audit every tiny model in a directory for feature-contract drift, the case where libvmaf's feature columns changed but a shipped model expects the old shape.

validate-norm¶

Compare a sidecar's declared feature normalisation against real data. It flags features whose declared mean drifts more than 1 sigma from the observed mean, or where more than 5% of samples lie more than 3 sigma from it. This catches silent normalisation breakage that makes inference under-predict.

profile¶

Measure latency and peak-RSS delta per (provider, shape): mean, p50 and p99. Use it to pick a deployment target or as a CI latency gate.

cross-backend¶

Run a model on the CPU provider and every other available ORT provider and diff the outputs, to catch provider-specific numerical drift.

audit-learned-filter¶

Pre-deploy audit of a learned_filter model. It runs the filter over a frame corpus and flags four failures: mean shift, std inflation, clipping at codec boundaries and SSIM collapse.

bisect-model-quality¶

Binary-search an ordered list of checkpoints (head assumed good, tail assumed bad) for the first one that violates a PLCC, SROCC or RMSE gate on a held-out feature cache. If the list is not ordered good to bad, it exits 0 with the verdict "no regression detected" or "first model already fails". Companion to the /bisect-model-quality skill.

Register¶

register¶

Write the sidecar metadata JSON (<model>.json) beside a shipped ONNX model (ADR-0211).

JSON report provenance¶

Every vmaf-train subcommand that writes a durable JSON report with --json adds an ADR-0661 run_provenance block. It covers validate-norm, profile, audit-learned-filter, quantize-int8, cross-backend and bisect-model-quality.

The block records:

• entrypoint: ai/src/vmaf_train/cli.py plus a SHA-256 of the CLI file.
• argv and args: the invoked command arguments and parsed option values.
• inputs: the model, feature table, calibration table, frame corpus or model list behind the report.
• outputs: the JSON report path, plus generated model outputs where the command writes one, such as quantize-int8 --output.

Use the block when attaching reports to model cards, promotion PRs or regression investigations: it is the reproducibility pointer to the exact files and thresholds.

Related¶

• vmaf-tune — encode automation that consumes the models produced here.
• vmaf — scoring CLI (--tiny-model flag).
• docs/ai/training.md — training-corpus pipeline.
• docs/ai/quantization.md — post-training quantisation (ADR-0173).