NVDINOv2
NVDINOv2 for self-supervised visual representation learning. Trains vision transformers via self-distillation (teacher-student) without labels. Produces general-purpose visual features.
Set train.pretrained_model_path for pretrained ViT weights.
For TAO Deploy TensorRT actions (
), read
references/tao-deploy-nvdinov2.md
first. Deploy spec templates live in this skill's
folder with the
spec_template_deploy_*.yaml
prefix.
Dataclass Schemas
Generated TAO Core schemas are packaged in
schemas/<action>.schema.json
, with
listing available actions. Each generated schema also emits
references/spec_template_<action>.yaml
from the schema top-level
field. AutoML enablement is declared at the model layer in
references/skill_info.yaml
via
. Runnable AutoML still requires
schemas/train.schema.json
and
references/spec_template_train.yaml
to exist and parse. Use the packaged train schema for
automl_default_parameters
,
automl_disabled_parameters
, defaults, min/max bounds, enums, option weights, math conditions, dependencies, and popular parameters. Do not expect
at runtime; maintainers regenerate schemas/templates before packaging the skill bank.
Train Action Policy
This model is AutoML-enabled at the model layer. Before handling any train-stage request, read
references/skill_info.yaml
and resolve the run override from either an explicit
value or the user's workflow request. Treat phrases like "turn off AutoML", "disable AutoML", "no HPO", or "plain training" as
for this run only; otherwise default to
. When
,
, and both
schemas/train.schema.json
and
references/spec_template_train.yaml
are packaged, route the train action through
tao-skill-bank:tao-run-automl
by default with this model's
. Preserve workflow/application overrides for datasets, specs, output directories, GPU/platform settings, parent checkpoints, and
. Use direct model training only when
or the packaged train schema/template is missing; in the missing-schema case, report that AutoML is enabled but not runnable for this model until schemas are generated.
Non-train actions such as
,
,
, and deploy flows stay in this model skill. The per-run
override does not change model metadata.
Training Requirements
- Dataset type: image_classification
- Formats: ssl
- Monitoring metric: train_loss
Per-Action Dataset Requirements
| Action | Spec Key | Source | Files | List? |
|---|
| distill | dataset.train_dataset.images_dir | train_datasets | images_train.tar.gz | No |
| inference | dataset.test_dataset.images_dir | inference_dataset | images_test.tar.gz | No |
| train | dataset.train_dataset.images_dir | train_datasets | images_train.tar.gz | No |
Typical Spec Overrides
Data source overrides are
mandatory for every action — the agent MUST construct data source paths from the Per-Action Dataset Requirements table above and include them in
.
python
S3_TRAIN = "s3://bucket/data/train"
S3_EVAL = "s3://bucket/data/eval"
train (mandatory data sources):
python
{
"train.num_gpus": 1,
"train.num_epochs": 10,
"train.checkpoint_interval": 10,
"dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}
distill (mandatory data sources):
python
{
"dataset.train_dataset.images_dir": f"{S3_TRAIN}/images_train.tar.gz",
}
inference (mandatory data sources):
python
{
"dataset.test_dataset.images_dir": f"{S3_EVAL}/images_test.tar.gz",
}
Eval Dataset
Optional. SSL training does not use labels. Evaluation is downstream task-specific.
Important Parameters
- model.backbone.teacher_type: Teacher ViT variant. Default vit_l (ViT-Large).
- model.backbone.student_type: Student ViT variant. Default vit_l. Typically matches teacher.
- model.backbone.img_size: Input image size. Default 518. Higher resolution produces better features but costs more memory.
- model.backbone.patch_size: ViT patch size. Default 14.
- dataset.batch_size: Per-GPU batch size. Default 4. SSL training is memory-intensive due to dual (teacher+student) forward passes.
- train.layerwise_decay: Layer-wise learning rate decay. Important for ViT fine-tuning.
- train.clip_grad_norm: Gradient clipping. Important for stable SSL training.
Multi-GPU / Multi-Node
Launch method: Lightning-managed (single
process, Lightning spawns workers).
| Spec Key | Description | Default |
|---|
| Number of GPUs | 1 |
| GPU device indices | [0] |
| Number of nodes | 1 |
- Strategy: (Lightning picks best strategy automatically)
- is always enabled — critical for SSL training with teacher-student framework
- Multi-GPU strongly recommended (4-8 GPUs) for meaningful SSL training
Multi-node env vars (set by orchestrator):
,
,
,
,
.
Hardware
Minimum 4 GPU(s), recommended 8 GPU(s). 40GB+ (A100 recommended) VRAM per GPU. SSL with ViT-Large teacher+student is very memory-intensive. Requires A100 40GB+ GPUs. Multi-GPU strongly recommended.
Error Patterns
CUDA out of memory: ViT-Large teacher+student with img_size=518 requires 40GB+ GPU memory. Reduce batch_size, img_size, or use smaller ViT variant.
Slow convergence: SSL needs many epochs. Default 10 is for quick testing; production runs typically use 100+ epochs.
Spec Param / Parent Model Inference
Model-specific inference mappings belong in this MD file, not in
. Generated runners should read this section and apply the mappings with SDK helpers before
. This mirrors the old microservices
flow.
Inference mappings from TAO Core
:
| Action | Spec Field | Inference Function | Meaning |
|---|
| distill | | | encryption key |
| distill | model.distill.pretrained_non_distill_pl_model_path
| | model file inferred from the parent job results folder |
| distill | | | current job results directory |
| export | | | encryption key |
| export | | | model file inferred from the parent job results folder |
| export | | | output ONNX path |
| export | | | current job results directory |
| inference | | | encryption key |
| inference | | | model file inferred from the parent job results folder |
| inference | | | model file inferred from the parent job results folder |
| inference | | | current job results directory |
| train | | | encryption key |
| train | | | current job results directory |
| train | train.pretrained_model_path
| | PTM when no resume checkpoint exists |
| train | train.resume_training_checkpoint_path
| | model file inferred from the current job results folder |
For
or
, pass the upstream train/export/AutoML child job id as
. The SDK lists the parent result folder, filters checkpoint artifacts, and returns the selected model file or folder. Do not add these mappings back to
and do not patch generated runner scripts to guess checkpoint paths.
Deployment
- tao-deploy-nvdinov2 — NvDINOv2 deploy workflow for TensorRT engine generation using TAO Deploy.