Start with Use TrainCheck if you want the full workflow. This page explains the traincheck-collect command.
traincheck-collect instruments a PyTorch training script and writes trace files. Use it for two jobs:
- Full reference collection for invariant inference.
- Selective target collection for checking with an existing invariant file.
Use full collection on a known-good run:
traincheck-collect \
--pyscript train.py \
--models-to-track model \
--output-dir reference_traceThis command runs train.py, tracks the Python variable named model, and writes trace files into reference_trace/.
Use selective collection when you already have an invariant file:
traincheck-collect \
--pyscript train.py \
--models-to-track model \
--invariants invariants.json \
--output-dir target_trace--invariants tells TrainCheck which APIs and variables matter for checking. This usually reduces target-run overhead compared with full reference collection.
Do not combine --invariants with --use-full-instr when you want selective collection.
Sampling is also configured on traincheck-collect:
traincheck-collect \
--pyscript train.py \
--models-to-track model \
--invariants invariants.json \
--sampling-interval 10 \
--warm-up-steps 10 \
--output-dir target_traceThis traces the warm-up steps, then traces every tenth step. Use sampling for long target runs after you have confirmed TrainCheck works on a short run.
Use --use-config when the collection command needs repeated options:
traincheck-collect --use-config --config traincheck.ymlExample:
pyscript: ./train.py
shscript: ./run.sh
modules_to_instr:
- torch
models_to_track:
- model
model_tracker_style: subclass
copy_all_files: false
output_dir: traincheck_traceConfig keys use underscores, not hyphens. For example, the CLI flag --output-dir becomes output_dir in YAML.
--pyscript: Python entry point for the training program.--shscript: shell script used to launch the Python program.--models-to-track: model variable names to track.--modules-to-instr: Python modules to instrument, usuallytorch.--invariants: invariant files for selective collection.--output-dir: directory for traces and logs.--sampling-interval: collect every Nth step after warm-up.--warm-up-steps: collect the first N steps.--copy-all-files: copy files beside the training script into the output directory.--model-tracker-style: choosesubclass(default),proxy, orsampler.
The default subclass tracker keeps the model as a PyTorch nn.Module and
tracks its parameters using a torch.nn.Parameter subclass. It requires
--models-to-track (or models_to_track in YAML); choosing a tracker alone does
not enable model tracking. Model tracking must be installed before creating the
optimizer, so the optimizer uses the tracked parameters.
The proxy tracker remains available through --model-tracker-style proxy or
model_tracker_style: proxy in YAML. It wraps model objects and dynamically
accessed attributes, so its coverage differs from the parameter-focused subclass
tracker. Explicit CLI/YAML selections remain supported, and saved
instrumentation options are not rewritten.
For invariants inferred from proxy traces, keep proxy explicit when collecting
the checking trace. To migrate, collect a new reference trace with subclass,
infer new invariants, and verify that the intended rules actually activate during
checking. Zero violations alone does not establish coverage; module or buffer
rules from a proxy trace may not activate with parameter-only tracking.
Selective collection keeps its existing policy: invariants requiring immediate
variable tracking use the selected subclass or proxy tracker. When only
cross-variable consistency rules require model tracking, collection uses the
sampler tracker to observe parameters after optimizer steps. If no invariant
requires model tracking, selective collection disables it.
Run the command help for the complete option list:
traincheck-collect --helpThe output directory contains trace files, environment metadata, logs, and the instrumented training script. The checker accepts the full output directory through -f or --trace-folders.
traincheck-check -f target_trace -i invariants.json