Tao Train Image Classification

NVIDIAOfficialCreated by NVIDIA Corporation

Skill
Tool integration
Data and AI
Apache-2.0
Medium risk

From the author

PyTorch-based TAO image classification. Supports a wide range of backbones (FAN, EfficientNet, ResNet, etc.) with distillation and quantization for deployment. Use when training, evaluating, distilling, quantizing, exporting, or running inference for a TAO image-classification (PyT) model. Trigger phrases include "train image classifier", "TAO classification", "ResNet/EfficientNet/FAN backbone classifier", "classification-pyt".

— the description by NVIDIA Corporation

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

This skill provides a PyTorch-based framework for image classification tasks within the NVIDIA TAO ecosystem. It supports a variety of model backbones, including FAN, EfficientNet, and ResNet. The skill facilitates the full model lifecycle, including training, evaluation, distillation, quantization, and export. It includes AutoML capabilities for training and provides integration paths for TensorRT deployment.

What it does

Trains image classification models using various supported backbones.
Supports model distillation and quantization for optimized deployment.
Exports models to ONNX format for further processing.
Integrates with TAO Deploy for TensorRT engine generation and inference.
Provides AutoML-enabled training workflows.
Handles evaluation of model checkpoints and TensorRT engines.

When to use it

When you need to train an image classification model using PyTorch backbones.
When you need to optimize a model through distillation or quantization.
When you need to prepare a model for deployment via TensorRT.
When you need to perform inference or evaluation on classification datasets.

Advantages

Supports a wide range of model backbones.
Includes built-in support for model distillation and quantization.
Automates training workflows with AutoML integration.
Provides clear paths for TensorRT deployment and engine generation.

Usage examples

“Train an image classifier using the ResNet backbone.”: The agent initiates the training process using the specified ResNet backbone, applying the required dataset and configuration overrides.
“Quantize my trained classification model.”: The agent runs the quantization action, utilizing the provided calibration dataset to optimize the model for deployment.
“Generate a TensorRT engine for my exported ONNX model.”: The agent triggers the gen_trt_engine action to convert the ONNX model into a TensorRT engine for optimized inference.
“Evaluate my model performance on the test dataset.”: The agent runs the evaluation workflow, calculating metrics based on the provided test images and model checkpoint.

Install

Download the skill as a ZIP:
https://autogineer.com/agent-skills/nvidia/tao-train-image-classification/download

Unzip it into the cross-client skills folder of your project:

mkdir -p .agents/skills && unzip tao-train-image-classification.zip -d .agents/skills

Documentation by NVIDIA Corporation

Classification PyT

Standalone install? If this session was not initialized by the TAO skill bank plugin, run the tao-setup skill first (host preflight, credentials, cross-skill discovery).

PyTorch image classification. Supports a wide range of backbones (FAN, EfficientNet, ResNet, etc.) with distillation and quantization for deployment.

Set model.backbone.pretrained_backbone_path for backbone weights or train.pretrained_model_path for full model.

For TAO Deploy TensorRT actions (gen_trt_engine, TensorRT evaluate, and TensorRT inference), read references/tao-deploy-image-classification.md first. Deploy spec templates live in this skill's references/ folder with the spec_template_deploy_*.yaml prefix.

Dataclass Schemas

Generated TAO Core schemas are packaged in schemas/<action>.schema.json, with schemas/manifest.json listing available actions. Each generated schema also emits references/spec_template_<action>.yaml from the schema top-level default field. AutoML enablement is declared at the model layer in references/skill_info.yaml via automl_enabled. Runnable AutoML for an action requires schemas/<action>.schema.json and references/spec_template_<action>.yaml to exist and parse. Use the packaged selected-action schema for automl_default_parameters, automl_disabled_parameters, defaults, min/max bounds, enums, option weights, math conditions, dependencies, and popular parameters. Do not expect ~/tao-core 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 automl_policy value or the user's workflow request. Use automl_policy: on by default and only expose on / off in new launch prompts. Treat phrases like "turn off AutoML", "disable AutoML", "no HPO", or "plain training" as automl_policy: off for this run only. When automl_policy: on, automl_enabled: true, 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 skill_dir. Preserve workflow/application overrides for datasets, specs, output directories, GPU/platform settings, parent checkpoints, and automl_policy. Use direct model training only when automl_policy: off 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 evaluate, inference, export, and deploy flows stay in this model skill. The per-run automl_policy override does not change model metadata.

Training Requirements

Dataset type: image_classification
Formats: classification_pyt
Monitoring metric: val_acc_1
AutoML train scope: exclude distill.teacher.* parameters from ordinary train searches. Those fields belong to the separate distill action and do not change classification_pyt train behavior even though they appear in the combined generated schema.

Per-Action Dataset Requirements

Action	Spec Key	Source	Files	List?
distill	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	No
distill	dataset.classes_file	train_datasets	classes.txt	No
distill	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	No
evaluate	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	No
evaluate	dataset.classes_file	eval_dataset	classes.txt	No
evaluate	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz	No
export	dataset.root_dir	train_datasets		No
inference	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	No
inference	dataset.classes_file	eval_dataset	classes.txt	No
inference	dataset.test_dataset.images_dir	inference_dataset	images_test.tar.gz	No
quantize	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	No
quantize	dataset.classes_file	train_datasets	classes.txt	No
quantize	dataset.val_dataset.images_dir	eval_dataset	images_val.tar.gz	No
quantize	dataset.quant_calibration_dataset.images_dir	calibration_dataset	images_train.tar.gz	No
train	dataset.train_dataset.images_dir	train_datasets	images_train.tar.gz	No
train	dataset.classes_file	train_datasets	classes.txt	No
train	dataset.val_dataset.images_dir	eval_dataset	images_val.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 spec_overrides.

TRAIN_IMAGES_DIR = "/workspace/data/extracted/train/images_train"
VAL_IMAGES_DIR = "/workspace/data/extracted/val/images_val"
TEST_IMAGES_DIR = "/workspace/data/extracted/test/images_test"
CLASSES_FILE = "/workspace/data/s3/classes.txt"
WRITABLE_RESULTS_DIR = "/results"   # must be a writable bind, not the image CWD

For local Docker, download the S3 archives, extract them first, and point dataset.*.images_dir at the extracted class-root folder. Do not pass images_train.tar.gz, images_val.tar.gz, or images_test.tar.gz directly to local Docker specs; the skill metadata declares these inputs as folders.

train (mandatory data sources):

{
    "train.num_epochs": 2,
    "train.validation_interval": 2,
    "train.checkpoint_interval": 2,
    "train.num_gpus": 1,
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.root_dir": WRITABLE_RESULTS_DIR,
}

dataset.root_dir is mandatory for train, not just for export. CLDataset writes a classes.txt into root_dir during setup, so it must point at a writable bind mount (the results mount is the natural choice). The spec template defaults it to '', which makes the write land in the container working directory — that succeeds under a root container but fails under any runner that drops to a non-root user:

PermissionError: [Errno 13] Permission denied: 'classes.txt'
  .../classification_pyt/dataloader/dataset.py, in CLDataset.__init__

The local Docker launcher uses the host UID:GID whenever a writable results bind is present, so non-root training hits this unless dataset.root_dir is set.

export (mandatory data sources):

{
    "export.input_height": 224,
    "export.input_width": 224,
    "dataset.root_dir": "/workspace/data/extracted",
}

gen_trt_engine:

{
    "gen_trt_engine.tensorrt.data_type": "fp16",
}

inference (mandatory data sources):

{
    "dataset.batch_size": 1,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.test_dataset.images_dir": TEST_IMAGES_DIR,
}

distill (mandatory data sources):

{
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "train.optim.policy": "step",
}

evaluate (mandatory data sources):

{
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.test_dataset.images_dir": TEST_IMAGES_DIR,
}

quantize (mandatory data sources):

{
    "dataset.train_dataset.images_dir": TRAIN_IMAGES_DIR,
    "dataset.classes_file": CLASSES_FILE,
    "dataset.val_dataset.images_dir": VAL_IMAGES_DIR,
    "dataset.quant_calibration_dataset.images_dir": TRAIN_IMAGES_DIR,
}

Eval Dataset

Optional. Validation images are provided as a separate tar alongside training images. For small smoke datasets that do not provide a separate images_test.tar.gz, set dataset.test_dataset.images_dir to the validation archive so evaluate and inference still exercise the checkpoint handoff.

Important Parameters

dataset.num_classes: Number of classes. Default 20. Must match the number of subdirectories in your image tarballs.
model.backbone.type: Default fan_small_12_p4_hybrid. Supported backbones and their head in_channels (from model_params_mapping.py): FAN: fan_tiny, fan_small_12_p4_hybrid, fan_base_16_p4_hybrid, fan_large_16_p4_hybrid. GCViT: gcvit_tiny through gcvit_large. FasterViT: fastervit_0 through fastervit_6. ViT/EVA/DINO: vit_large_patch14_dinov2, eva02_large_patch14, etc. SigLIP-CLIPA: ViT-H-14-SigLIP-CLIPA-224, etc. Some backbones require non-default input resolution (384, 512, 768).
dataset.classes_file: Path to classes.txt listing class names.
train.optim.lr: Learning rate. Default 6e-5.
dataset.img_size: Input image size. Default 224.
dataset.batch_size: Per-GPU batch size. Default 8.

Multi-GPU / Multi-Node

Launch method: Lightning-managed (single python process, Lightning spawns workers).

Spec Key	Description	Default
`train.num_gpus`	Number of GPUs	1
`train.gpu_ids`	GPU device indices	[0]
`train.num_nodes`	Number of nodes	1

Multi-GPU strategy: ddp_find_unused_parameters_true
No fsdp support

Multi-node env vars (set by orchestrator): WORLD_SIZE, NODE_RANK, MASTER_ADDR, MASTER_PORT, NUM_GPU_PER_NODE.

Hardware

Minimum 1 GPU(s), recommended 2 GPU(s). 16GB+ (V100 or A100) VRAM per GPU. Classification is generally lightweight. Most backbones at 224x224 fit well on 16GB GPUs with batch_size=8.

Error Patterns

CUDA out of memory: Reduce batch_size or use a smaller backbone.

num_classes mismatch: Ensure dataset.num_classes matches the actual class directories in your image tarballs and classes.txt.

Empty class directory: Every class in classes.txt must have at least one image in the corresponding subdirectory.

Distill scheduler default: The bundled distill template and schema use train.optim.policy: step. Keep that setting for distill specs unless the container implementation is updated; the 7.0 PyT distiller does not assign a scheduler interval for train.optim.policy: linear.

Checkpoint handoff: Training produces model_epoch_*.pth checkpoints and a classifier_model_latest.pth symlink. For evaluate, inference, export, quantize, distill, and resume, select the exact intended epoch checkpoint through the SDK resolver; use the latest symlink only when the user explicitly requests latest.

Spec Param / Parent Model Inference

Model-specific inference mappings belong in this MD file, not in config.json. Generated runners should read this section and apply the mappings with SDK helpers before create_job(). This mirrors the old microservices infer_params.py flow.

Inference mappings from TAO Core classification_pyt.config.json:

Action	Spec Field	Inference Function	Meaning
distill	`distill.pretrained_teacher_model_path`	`parent_model`	model file inferred from the parent job results folder
distill	`results_dir`	`output_dir`	current job results directory
evaluate	`evaluate.checkpoint`	`parent_model`	model file inferred from the parent job results folder
evaluate	`results_dir`	`output_dir`	current job results directory
export	`export.checkpoint`	`parent_model`	model file inferred from the parent job results folder
export	`export.onnx_file`	`create_onnx_file`	output ONNX path
export	`results_dir`	`output_dir`	current job results directory
gen_trt_engine	`gen_trt_engine.onnx_file`	`parent_model`	model file inferred from the parent job results folder
gen_trt_engine	`gen_trt_engine.trt_engine`	`create_engine_file`	output TensorRT engine path
gen_trt_engine	`results_dir`	`output_dir`	current job results directory
inference	`inference.checkpoint`	`parent_model`	model file inferred from the parent job results folder
inference	`inference.trt_engine`	`parent_model`	model file inferred from the parent job results folder
inference	`results_dir`	`output_dir`	current job results directory
quantize	`quantize.model_path`	`parent_model`	model file inferred from the parent job results folder
quantize	`results_dir`	`output_dir`	current job results directory
train	`model.backbone.pretrained_backbone_path`	`ptm_if_no_resume_model`	PTM when no resume checkpoint exists
train	`results_dir`	`output_dir`	current job results directory
train	`train.pretrained_model_path`	`ptm_if_no_resume_model`	PTM when no resume checkpoint exists
train	`train.resume_training_checkpoint_path`	`resume_model`	model file inferred from the current job results folder

For parent_model or parent_model_folder, pass the upstream train/export/AutoML child job id as parent_job_id. 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 config.json and do not patch generated runner scripts to guess checkpoint paths.

Deployment

tao-deploy-image-classification

Files

ATTRIBUTION.mdAdded by AutoGineer960 B
BENCHMARK.md7.1 KiB
LICENSELicence file added by AutoGineer11.2 KiB
SKILL.md13.9 KiB
config/skillspector-baseline.yaml834 B
evals/evals.json823 B
references/skill_info.yaml7.0 KiB
references/spec_template_deploy_evaluate.yaml200 B
references/spec_template_deploy_gen_trt_engine.yaml217 B
references/spec_template_deploy_inference.yaml168 B
references/spec_template_distill.yaml2.9 KiB
references/spec_template_evaluate.yaml3.1 KiB
references/spec_template_export.yaml3.2 KiB
references/spec_template_gen_trt_engine.yaml3.3 KiB
references/spec_template_inference.yaml3.1 KiB
references/spec_template_quantize.yaml2.9 KiB
references/spec_template_train.yaml2.9 KiB
references/tao-deploy-image-classification.md5.8 KiB
references/tao-deploy-image-classification.skill_info.yaml2.4 KiB
schemas/distill.schema.json61.9 KiB
schemas/evaluate.schema.json65.0 KiB
schemas/export.schema.json65.9 KiB
schemas/gen_trt_engine.schema.json69.8 KiB
schemas/inference.schema.json65.0 KiB
schemas/manifest.json28.4 KiB
schemas/quantize.schema.json61.9 KiB
schemas/train.schema.json61.9 KiB
skill-card.md4.2 KiB
skill.oms.sig9.9 KiB

schemas/train.schema.json

{
  "automl_default_parameters": [
    "dataset.augmentation.random_color.hue",
    "dataset.batch_size",
    "dataset.augmentation.random_erase.enable",
    "train.optim.weight_decay",
    "dataset.augmentation.random_aug.enable",
    "train.optim.betas",
    "dataset.workers",
    "train.optim.momentum",
    "dataset.augmentation.random_flip.vflip_probability",
    "dataset.augmentation.random_color.enable",
    "dataset.augmentation.random_rotate.rotate_probability",
    "model.backbone.freeze_backbone",
    "distill.teacher.backbone.freeze_norm",
    "dataset.augmentation.random_color.saturation",
    "dataset.augmentation.random_color.contrast",
    "model.backbone.freeze_norm",
    "train.optim.lr",
    "dataset.augmentation.random_color.color_probability",
    "dataset.augmentation.random_rotate.enable",
    "dataset.augmentation.with_scale_random_crop.enable",
    "dataset.augmentation.random_color.brightness",
    "distill.teacher.backbone.freeze_backbone",
    "dataset.augmentation.random_flip.enable",
    "dataset.augmentation.random_erase.erase_probability",
    "dataset.augmentation.random_flip.hflip_probability"
  ],
  "automl_disabled_parameters": [
    "quantize",
    "dataset.augmentation.with_scale_random_crop.scale_range",
    "train.cudnn",
    "gen_trt_engine.tensorrt.calibration.cal_image_dir",
    "distill.teacher.head.custom_args",
    "quantize.backend_kwargs",
    "train.gpu_ids",
    "wandb.tags",
    "model.backbone",
    "distill.teacher.backbone",
    "distill.teacher.head.loss",
    "quantize.skip_names",
    "train.tensorboard",
    "dataset.augmentation.random_rotate",
    "dataset.train_dataset",
    "distill.teacher.head.topk",
    "evaluate",
    "inference",
    "train",
    "distill",
    "dataset.augmentation",
    "dataset.augmentation.random_erase",
    "dataset.test_dataset",
    "gen_trt_engine",
    "gen_trt_engine.tensorrt.layers_precision",
    "dataset",
    "gen_trt_engine.tensorrt",
    "train.optim.policy_params",
    "dataset.val_dataset",
    "quantize.layers",
    "dataset.quant_calibration_dataset",
    "model.head",
    "dataset.augmentation.random_color",
    "model",
    "model.head.topk",
    "model.head.custom_args",
    "evaluate.gpu_ids",
    "train.optim",
    "dataset.augmentation.with_scale_random_crop",
    "distill.teacher",
    "gen_trt_engine.tensorrt.calibration",
    "distill.teacher.head",
    "dataset.augmentation.random_aug",
    "train.optim.skip_names",
    "dataset.train_nolabel",
    "dataset.augmentation.std",
    "export",
    "dataset.augmentation.random_rotate.angle_list",
    "wandb",
    "inference.gpu_ids",
    "dataset.augmentation.random_flip",
    "model.head.loss",
    "dataset.augmentation.mean"
  ],
  "default": {
    "dataset": {
      "augmentation": {
        "mean": [
          0.485,
          0.456,
          0.406
        ],
        "mixup_alpha": 0.4,
        "mixup_cutmix": false,
        "random_aug": {
          "enable": true
        },
        "random_color": {
          "brightness": 0.3,
          "color_probability": 0.5,
          "contrast": 0.3,
          "enable": true,
          "hue": 0.0,
          "saturation": 0.3
        },
        "random_erase": {
          "enable": true,
          "erase_probability": 0.2
        },
        "random_flip": {
          "enable": true,
          "hflip_probability": 0.5,
          "vflip_probability": 0.5
        },
        "random_rotate": {
          "angle_list": [
            90,
            180,
            270
          ],
          "enable": true,
          "rotate_probability": 0.5
        },
        "std": [
          0.229,
          0.224,
          0.225
        ],
        "with_random_blur": true,
        "with_random_crop": true,
        "with_scale_random_crop": {
          "enable": true,
          "scale_range": [
            1,
            1.2
          ]
        }
      },
      "batch_size": 8,
      "classes_file": "",
      "dataset": "CLDataset",
      "img_size": 224,
      "num_classes": 20,
      "quant_calibration_dataset": {
        "images_dir": ""
      },
      "root_dir": "",
      "shuffle": true,
      "test_dataset": {
        "images_dir": ""
      },
      "train_dataset": {
        "images_dir": ""
      },
      "train_nolabel": {
        "folder_path": ""
      },
      "val_dataset": {
        "images_dir": ""
      },
      "workers": 1
    },
    "distill": {
      "loss_lambda": 0.5,
      "loss_type": "KL",
      "mlp_hidden_size": 1024,
      "mlp_num_inner": 0,
      "mode": "auto",
      "pretrained_teacher_model_path": "???",
      "results_dir": "",
      "teacher": {
        "backbone": {
          "freeze_backbone": false,
          "freeze_norm": false,
          "pretrained_backbone_path": "",
          "type": "fan_small_12_p4_hybrid"
        },
        "head": {
          "binary": false,
          "in_channels": 448,
          "loss": {
            "label_smooth_val": 0.0,
            "type": "CrossEntropyLoss"
          },
          "topk": [
            1
          ],
          "type": "TAOLinearClsHead"
        }
      },
      "use_mlp": true
    },
    "encryption_key": "",
    "model": {
      "backbone": {
        "freeze_backbone": false,
        "freeze_norm": false,
        "pretrained_backbone_path": "",
        "type": "fan_small_12_p4_hybrid"
      },
      "head": {
        "binary": false,
        "in_channels": 448,
        "loss": {
          "label_smooth_val": 0.0,
          "type": "CrossEntropyLoss"
        },
        "topk": [
          1
        ],
        "type": "TAOLinearClsHead"
      }
    },
    "model_name": "",
    "quantize": {
      "algorithm": "minmax",
      "backend": "torchao",
      "backend_kwargs": {},
      "device": "cuda",
      "layers": [],
      "mode": "weight_only_ptq",
      "model_path": "",
      "results_dir": "",
      "skip_names": []
    },
    "results_dir": "",
    "train": {
      "checkpoint_interval": 1,
      "checkpoint_interval_unit": "epoch",
      "clip_grad_norm": 2.0,
      "cudnn": {
        "benchmark": false,
        "deterministic": true
      },
      "ema_decay": 0.998,
      "enable_ema": false,
      "gpu_ids": [
        0
      ],
      "num_epochs": 10,
      "num_gpus": 1,
      "num_nodes": 1,
      "optim": {
        "betas": [
          0.9,
          0.999
        ],
        "lr": 6e-05,
        "momentum": 0.9,
        "monitor_name": "val_loss",
        "optim": "adamw",
        "policy": "linear",
        "policy_params": {
          "gamma": 0.1,
          "step_size": 30
        },
        "skip_names": [],
        "warmup_epochs": 0,
        "weight_decay": 0.01
      },
      "precision": "fp32",
      "pretrained_model_path": "",
      "results_dir": "",
      "resume_training_checkpoint_path": "",
      "seed": 1234,
      "tensorboard": {
        "enabled": false,
        "infrequent_logging_frequency": 2
      },
      "validation_interval": 1
    },
    "wandb": {
      "enable": true,
      "entity": "",
      "group": "",
      "name": "TAO Toolkit Training",
      "project": "TAO Toolkit",
      "reinit": false,
      "run_id": "",
      "save_code": false,
      "sync_tensorboard": false,
      "tags": [
        "tao-toolkit"
      ]
    }
  },
  "popular": {
    "evaluate": {
      "gpu_ids": [
        0
      ],
      "num_gpus": 1,
      "num_nodes": 1
    },
    "gen_trt_engine": {
      "batch_size": -1,
      "gpu_id": 0,
      "tensorrt": {
        "calibration": {
          "cal_batch_size": 1,
          "cal_batches": 1
        },
        "max_batch_size": 1,
        "min_batch_size": 1,
        "opt_batch_size": 1
      }
    },
    "inference": {
      "gpu_ids": [
        0
      ],
      "num_gpus": 1,
      "num_nodes": 1
    },
    "train": {
      "checkpoint_interval": 1,
      "gpu_ids": [
        0
      ],
      "num_epochs": 10,
      "num_gpus": 1,
      "num_nodes": 1,
      "validation_interval": 1
    }
  },
  "properties": {
    "automl_disabled_parameters": [
      "wandb",
      "model",
      "dataset",
      "train",
      "evaluate",
      "inference",
      "export",
      "gen_trt_engine",
      "distill",
      "quantize"
    ],
    "dataset": {
      "automl_default_parameters": [
        "dataset.batch_size",
        "dataset.workers"
      ],
      "automl_disabled_parameters": [
        "dataset.augmentation",
        "dataset.train_dataset",
        "dataset.train_nolabel",
        "dataset.val_dataset",
        "dataset.test_dataset",
        "dataset.quant_calibration_dataset"
      ],
      "automl_enabled": false,
      "default": {
        "augmentation": {
          "mean": [
            0.485,
            0.456,
            0.406
          ],
          "mixup_alpha": 0.4,
          "mixup_cutmix": false,
          "random_aug": {
            "enable": true
          },
          "random_color": {
            "brightness": 0.3,
            "color_probability": 0.5,
            "contrast": 0.3,
            "enable": true,
            "hue": 0.0,
            "saturation": 0.3
          },
          "random_erase": {
            "enable": true,
            "erase_probability": 0.2
          },
          "random_flip": {
            "enable": true,
            "hflip_probability": 0.5,
            "vflip_probability": 0.5
          },
          "random_rotate": {
            "angle_list": [
              90,
              180,
              270
            ],
            "enable": true,
            "rotate_probability": 0.5
          },
          "std": [
            0.229,
            0.224,
            0.225
          ],
          "with_random_blur": true,
          "with_random_crop": true,
          "with_scale_random_crop": {
            "enable": true,
            "scale_range": [
              1,
              1.2
            ]
          }
        },
        "batch_size": 8,
        "classes_file": "",
        "dataset": "CLDataset",
        "img_size": 224,
        "num_classes": 20,
        "quant_calibration_dataset": {
          "images_dir": ""
        },
        "root_dir": "",
        "shuffle": true,
        "test_dataset": {
          "images_dir": ""
        },
        "train_dataset": {
          "images_dir": ""
        },
        "train_nolabel": {
          "folder_path": ""
        },
        "val_dataset": {
          "images_dir": ""
        },
        "workers": 1
      },
      "properties": {
        "augmentation": {
          "automl_disabled_parameters": [
            "dataset.augmentation.random_flip",
            "dataset.augmentation.random_rotate",
            "dataset.augmentation.random_color",
            "dataset.augmentation.random_erase",
            "dataset.augmentation.random_aug",
            "dataset.augmentation.with_scale_random_crop",
            "dataset.augmentation.mean",
            "dataset.augmentation.std"
          ],
          "automl_enabled": false,
          "default": {
            "mean": [
              0.485,
              0.456,
              0.406
            ],
            "mixup_alpha": 0.4,
            "mixup_cutmix": false,
            "random_aug": {
              "enable": true
            },
            "random_color": {
              "brightness": 0.3,
              "color_probability": 0.5,
              "contrast": 0.3,
              "enable": true,
              "hue": 0.0,
              "saturation": 0.3
            },
            "random_erase": {
              "enable": true,
              "erase_probability": 0.2
            },
            "random_flip": {
              "enable": true,
              "hflip_probability": 0.5,
              "vflip_probability": 0.5
            },
            "random_rotate": {
              "angle_list": [
                90,
                180,
                270
              ],
              "enable": true,
              "rotate_probability": 0.5
            },
            "std": [
              0.229,
              0.224,
              0.225
            ],
            "with_random_blur": true,
            "with_random_crop": true,
            "with_scale_random_crop": {
              "enable": true,
              "scale_range": [
                1,
                1.2
              ]
            }
          },
          "properties": {
            "mean": {
              "automl_enabled": false,
              "default": [
                0.485,
                0.456,
                0.406
              ],
              "description": "Mean for the augmentation",
              "title": "Mean",
              "type": "list"
            },
            "mixup_alpha": {
              "default": 0.4,
              "description": "Mixup alpha",
              "maximum": 1.0,
              "minimum": 0.0,
              "type": "float"
            },
            "mixup_cutmix": {
              "default": false,
              "description": "Flag to enable mixup and cutmix. Not recommended for binary classification.",
              "type": "bool"
            },
            "random_aug": {
              "automl_default_parameters": [
                "dataset.augmentation.random_aug.enable"
              ],
              "automl_enabled": false,
              "default": {
                "enable": true
              },
              "properties": {
                "enable": {
                  "automl_enabled": true,
                  "default": true,
                  "description": "Flag to enable Random Aug",
                  "type": "bool"
                }
              },
              "type": "collection"
            },
            "random_color": {
              "automl_default_parameters": [
                "dataset.augmentation.random_color.brightness",
                "dataset.augmentation.random_color.contrast",
                "dataset.augmentation.random_color.saturation",
                "dataset.augmentation.random_color.hue",
                "dataset.augmentation.random_color.enable",
                "dataset.augmentation.random_color.color_probability"
              ],
              "automl_enabled": false,
              "default": {
                "brightness": 0.3,
                "color_probability": 0.5,
                "contrast": 0.3,
                "enable": true,
                "hue": 0.0,
                "saturation": 0.3
              },
              "properties": {
                "brightness": {
                  "automl_enabled": true,
                  "default": 0.3,
                  "description": "Random Color Brightness",
                  "math_cond": "> 0.0",
                  "maximum": 2.0,
                  "minimum": 0.0,
                  "type": "float"
                },
                "color_probability": {
                  "automl_enabled": true,
                  "default": 0.5,
                  "description": "Random Color Probability",
                  "maximum": 1.0,
                  "minimum": 0.0,
                  "type": "float"
                },
                "contrast": {
                  "automl_enabled": true,
                  "default": 0.3,
                  "description": "Random Color Contrast",
                  "math_cond": "> 0.0",
                  "maximum": 2.0,
                  "minimum": 0.0,
                  "type": "float"
                },
                "enable": {
                  "automl_enabled": true,
                  "default": true,
                  "description": "Flag to enable Random Color",
                  "type": "bool"
                },
                "hue": {
                  "automl_enabled": true,
                  "default": 0.0,
                  "description": "Random Color Hue",
                  "math_cond": "> 0.0",
                  "maximum": 0.5,
                  "minimum": 0.0,
                  "type": "float"
                },
                "saturation": {
                  "automl_enabled": true,
                  "default": 0.3,
                  "description": "Random Color Saturation",
                  "math_cond": "> 0.0",
                  "maximum": 2.0,
                  "minimum": 0.0,
                  "type": "float"
                }
              },
              "type": "collection"
            },
            "random_erase": {
              "automl_default_parameters": [
                "dataset.augmentation.random_erase.enable",
                "dataset.augmentation.random_erase.erase_probability"
              ],
              "automl_enabled": false,
              "default": {
                "enable": true,
                "erase_probability": 0.2
              },
              "properties": {
                "enable": {
                  "automl_enabled": true,
                  "default": true,
                  "description": "Flag to enable Random Erase",
                  "type": "bool"
                },
                "erase_probability": {
                  "automl_enabled": true,
                  "default": 0.2,
                  "description": "Random Erase Probability",
                  "maximum": 1.0,
                  "minimum": 0.0,
                  "type": "float"
                }
              },
              "type": "collection"
            },
            "random_flip": {
              "automl_default_parameters": [
                "dataset.augmentation.random_flip.vflip_probability",
                "dataset.augmentation.random_flip.hflip_probability",
                "dataset.augmentation.random_flip.enable"
              ],
              "automl_enabled": false,
              "default": {
                "enable": true,
                "hflip_probability": 0.5,
                "vflip_probability": 0.5
              },
              "properties": {
                "enable": {
                  "automl_enabled": true,
                  "default": true,
                  "description": "Flag to enable augmentation",
                  "type": "bool"
                },
                "hflip_probability": {
                  "automl_enabled": true,
                  "default": 0.5,
                  "description": "Horizontal Flip probability",
                  "maximum": 1.0,
                  "minimum": 0.0,
                  "type": "float"
                },
                "vflip_probability": {
                  "automl_enabled": true,
                  "default": 0.5,
                  "description": "Vertical Flip probability",
                  "maximum": 1.0,
                  "minimum": 0.0,
                  "type": "float"
                }
              },
              "type": "collection"
            },
            "random_rotate": {
              "automl_default_parameters": [
                "dataset.augmentation.random_rotate.rotate_probability",
                "dataset.augmentation.random_rotate.enable"
              ],
              "automl_disabled_parameters": [
                "dataset.augmentation.random_rotate.angle_list"
              ],
              "automl_enabled": false,
              "default": {
                "angle_list": [
                  90,
                  180,
                  270
                ],
                "enable": true,
                "rotate_probability": 0.5
              },
              "properties": {
                "angle_list": {
                  "automl_enabled": false,
                  "default": [
                    90,
                    180,
                    270
                  ],
                  "description": "Random rotate angle probability",
                  "type": "list"
                },
                "enable": {
                  "automl_enabled": true,
                  "default": true,
                  "description": "Flag to enable augmentation",
                  "type": "bool"
                },
                "rotate_probability": {
                  "automl_enabled": true,
                  "default": 0.5,
                  "description": "Random Rotate probability",
                  "maximum": 1.0,
                  "minimum": 0.0,
                  "type": "float"
                }
              },
              "type": "collection"
            },
            "std": {
              "automl_enabled": false,
              "default": [
                0.229,
                0.224,
                0.225
              ],
              "description": "Standard deviation for the augmentation",
              "title": "Standard Deviation",
              "type": "list"
            },
            "with_random_blur": {
              "default": true,
              "description": "Flag to enable with_random_blur",
              "type": "bool"
            },
            "with_random_crop": {
              "default": true,
              "description": "Flag to enable with_random_crop",
              "type": "bool"
            },
            "with_scale_random_crop": {
              "automl_default_parameters": [
                "dataset.augmentation.with_scale_random_crop.enable"
              ],
              "automl_disabled_parameters": [
                "dataset.augmentation.with_scale_random_crop.scale_range"
              ],
              "automl_enabled": false,
              "default": {
                "enable": true,
                "scale_range": [
                  1,
                  1.2
                ]
              },
              "properties": {
                "enable": {
                  "automl_enabled": true,
                  "default": true,
                  "description": "Flag to enable Random Crop with Scale",
                  "type": "bool"
                },
                "scale_range": {
                  "automl_enabled": false,
                  "default": [
                    1,
                    1.2
                  ],
                  "description": "Random Scale range",
                  "type": "list"
                }
              },
              "type": "collection"
            }
          },
          "type": "collection"
        },
        "batch_size": {
          "automl_enabled": true,
          "default": 8,
          "description": "Batch size",
          "maximum": Infinity,
          "minimum": 1,
          "title": "Batch Size",
          "type": "int"
        },
        "classes_file": {
          "default": "",
          "description": "Path to the classes file",
          "type": "string"
        },
        "dataset": {
          "default": "CLDataset",
          "description": "dataset class",
          "enum": [
            "Dataset"
          ],
          "type": "categorical"
        },
        "img_size": {
          "default": 224,
          "description": "The input image size",
          "type": "int"
        },
        "num_classes": {
          "default": 20,
          "description": "The number of classes in the training data",
          "math_cond": ">=0",
          "maximum": Infinity,
          "minimum": 0,
          "type": "int"
        },
        "quant_calibration_dataset": {
          "automl_enabled": false,
          "default": {
            "images_dir": ""
          },
          "description": "Configuration for the quantization calibration dataset path",
          "properties": {
            "images_dir": {
              "default": "",
              "description": "Path to images directory for dataset",
              "title": "image directory",
              "type": "string"
            }
          },
          "title": "Quantization Calibration Dataset",
          "type": "collection"
        },
        "root_dir": {
          "default": "",
          "description": "Path to folder that contains classes.txt which indicate class name and train ID.         Can be optional then the mapping will be generated from pipeline.",
          "type": "string"
        },
        "shuffle": {
          "default": true,
          "description": "Shuffle dataloader",
          "type": "bool"
        },
        "test_dataset": {
          "automl_enabled": false,
          "default": {
            "images_dir": ""
          },
          "description": "Configuration for the testing dataset path",
          "properties": {
            "images_dir": {
              "default": "",
              "description": "Path to images directory for dataset",
              "title": "image directory",
              "type": "string"
            }
          },
          "title": "Testing Dataset",
          "type": "collection"
        },
        "train_dataset": {
          "automl_enabled": false,
          "default": {
            "images_dir": ""
          },
          "description": "Configuration for the training dataset path",
          "properties": {
            "images_dir": {
              "default": "",
              "description": "Path to images directory for dataset",
              "title": "image directory",
              "type": "string"
            }
          },
          "title": "Training Dataset",
          "type": "collection"
        },
        "train_nolabel": {
          "automl_enabled": false,
          "default": {
            "folder_path": ""
          },
          "properties": {
            "folder_path": {
              "default": "",
              "description": "Dataset directory path",
              "type": "string"
            }
          },
          "type": "collection"
        },
        "val_dataset": {
          "automl_enabled": false,
          "default": {
            "images_dir": ""
          },
          "description": "Configuration for the validation dataset path",
          "properties": {
            "images_dir": {
              "default": "",
              "description": "Path to images directory for dataset",
              "title": "image directory",
              "type": "string"
            }
          },
          "title": "Validation Dataset",
          "type": "collection"
        },
        "workers": {
          "automl_enabled": true,
          "default": 1,
          "description": "Workers",
          "maximum": Infinity,
          "minimum": 0,
          "title": "Workers",
          "type": "int"
        }
      },
      "type": "collection"
    },
    "distill": {
      "automl_disabled_parameters": [
        "distill.teacher"
      ],
      "automl_enabled": false,
      "default": {
        "loss_lambda": 0.5,
        "loss_type": "KL",
        "mlp_hidden_size": 1024,
        "mlp_num_inner": 0,
        "mode": "auto",
        "pretrained_teacher_model_path": "???",
        "results_dir": "",
        "teacher": {
          "backbone": {
            "freeze_backbone": false,
            "freeze_norm": false,
            "pretrained_backbone_path": "",
            "type": "fan_small_12_p4_hybrid"
          },
          "head": {
            "binary": false,
            "in_channels": 448,
            "loss": {
              "label_smooth_val": 0.0,
              "type": "CrossEntropyLoss"
            },
            "topk": [
              1
            ],
            "type": "TAOLinearClsHead"
          }
        },
        "use_mlp": true
      },
      "properties": {
        "loss_lambda": {
          "default": 0.5,
          "description": "The weight to be applied to the distillation loss as compared to task loss",
          "math_cond": "> 0.0 <= 1.0",
          "title": "distill weight",
          "type": "float"
        },
        "loss_type": {
          "default": "KL",
          "description": "Loss function for logits distillation.",
          "enum": [
            "\nKL(KLdivergence)",
            "\nCE(crossentropy)",
            "\nL1(L1loss)",
            "\nL2(L2loss)",
            "\nFD(smoothL1)",
            "\nCS(cosinesimilarity)",
            "\nBALANCED(balancedfeatureloss)",
            "\nMSE(meansquarederror)"
          ],
          "title": "Distillation loss type",
          "type": "categorical"
        },
        "mlp_hidden_size": {
          "default": 1024,
          "description": "MLP hidden size",
          "maximum": Infinity,
          "minimum": 0,
          "type": "int"
        },
        "mlp_num_inner": {
          "default": 0,
          "description": "MLP number of inner layers",
          "maximum": 10,
          "minimum": 0,
          "type": "int"
        },
        "mode": {
          "default": "auto",
          "description": "Distillation mode",
          "enum": [
            "logits",
            "summary",
            "spatial",
            "auto"
          ],
          "type": "categorical"
        },
        "pretrained_teacher_model_path": {
          "default": "???",
          "description": "Path to the pre-trained teacher model.",
          "title": "Pretrained teacher model path",
          "type": "string"
        },
        "results_dir": {
          "default": "",
          "description": "Path to where all the assets generated from a task are stored.",
          "title": "Results directory",
          "type": "string"
        },
        "teacher": {
          "automl_disabled_parameters": [
            "distill.teacher.backbone",
            "distill.teacher.head"
          ],
          "automl_enabled": false,
          "default": {
            "backbone": {
              "freeze_backbone": false,
              "freeze_norm": false,
              "pretrained_backbone_path": "",
              "type": "fan_small_12_p4_hybrid"
            },
            "head": {
              "binary": false,
              "in_channels": 448,
              "loss": {
                "label_smooth_val": 0.0,
                "type": "CrossEntropyLoss"
              },
              "topk": [
                1
              ],
              "type": "TAOLinearClsHead"
            }
          },
          "properties": {
            "backbone": {
              "automl_default_parameters": [
                "distill.teacher.backbone.freeze_backbone",
                "distill.teacher.backbone.freeze_norm"
              ],
              "automl_enabled": false,
              "default": {
                "freeze_backbone": false,
                "freeze_norm": false,
                "pretrained_backbone_path": "",
                "type": "fan_small_12_p4_hybrid"
              },
              "properties": {
                "freeze_backbone": {
                  "automl_enabled": true,
                  "default": false,
                  "description": "Flag to freeze backbone",
                  "type": "bool"
                },
                "freeze_norm": {
                  "automl_enabled": true,
                  "default": false,
                  "description": "Flag to freeze norm",
                  "type": "bool"
                },
                "pretrained_backbone_path": {
                  "default": "",
                  "description": "Path to the pretrained model",
                  "type": "string"
                },
                "type": {
                  "default": "fan_small_12_p4_hybrid",
                  "description": "Backbone architure",
                  "enum": [
                    "faster_vit_0_224",
                    "faster_vit_1_224",
                    "faster_vit_2_224",
                    "faster_vit_3_224",
                    "faster_vit_4_224",
                    "faster_vit_5_224",
                    "faster_vit_6_224",
                    "faster_vit_4_21k_224",
                    "faster_vit_4_21k_384",
                    "faster_vit_4_21k_512",
                    "faster_vit_4_21k_768",
                    "fan_tiny_12_p16_224",
                    "fan_small_12_p16_224_se_attn",
                    "fan_small_12_p16_224",
                    "fan_base_18_p16_224",
                    "fan_large_24_p16_224",
                    "fan_tiny_8_p4_hybrid",
                    "fan_small_12_p4_hybrid",
                    "fan_base_16_p4_hybrid",
                    "fan_large_16_p4_hybrid",
                    "fan_xlarge_16_p4_hybrid",
                    "fan_swin_tiny_patch4_window7_224",
                    "fan_swin_small_patch4_window7_224",
                    "fan_swin_base_patch4_window7_224",
                    "fan_swin_large_patch4_window7_224",
                    "vit_large_patch14_dinov2_swiglu",
                    "vit_large_patch14_dinov2_swiglu_legacy",
                    "vit_giant_patch14_reg4_dinov2_swiglu",
                    "efficientvit_b0",
                    "efficientvit_b1",
                    "efficientvit_b2",
                    "efficientvit_b3",
                    "efficientvit_l0",
                    "efficientvit_l1",
                    "efficientvit_l2",
                    "efficientvit_l3",
                    "vit_base_patch16",
                    "vit_large_patch16",
                    "vit_huge_patch14",
                    "convnext_tiny",
                    "convnext_small",
                    "convnext_base",
                    "convnext_large",
                    "convnext_xlarge",
                    "convnextv2_atto",
                    "convnextv2_femto",
                    "convnextv2_pico",
                    "convnextv2_nano",
                    "convnextv2_tiny",
                    "convnextv2_base",
                    "convnextv2_large",
                    "convnextv2_huge",
                    "hiera_tiny_224",
                    "hiera_small_224",
                    "hiera_base_224",
                    "hiera_base_plus_224",
                    "hiera_large_224",
                    "hiera_huge_224",
                    "resnet_18",
                    "resnet_34",
                    "resnet_50",
                    "resnet_101",
                    "resnet_152",
                    "resnet_18d",
                    "resnet_34d",
                    "resnet_50d",
                    "resnet_101d",
                    "resnet_152d",
                    "swin_tiny_patch4_window7_224",
                    "swin_small_patch4_window7_224",
                    "swin_base_patch4_window7_224",
                    "swin_large_patch4_window7_224",
                    "swin_base_patch4_window12_384",
                    "swin_large_patch4_window12_384",
                    "gc_vit_xxtiny",
                    "gc_vit_xtiny",
                    "gc_vit_tiny",
                    "gc_vit_small",
                    "gc_vit_base",
                    "gc_vit_large",
                    "gc_vit_base_384",
                    "gc_vit_large_384",
                    "edgenext_xx_small",
                    "edgenext_x_small",
                    "edgenext_small",
                    "edgenext_base",
                    "edgenext_xx_small_bn_hs",
                    "edgenext_x_small_bn_hs",
                    "edgenext_small_bn_hs",
                    "c_radio_p1_vit_huge_patch16_mlpnorm",
                    "c_radio_p2_vit_huge_patch16_mlpnorm",
                    "c_radio_p3_vit_huge_patch16_mlpnorm",
                    "c_radio_v2_vit_base_patch16",
                    "c_radio_v2_vit_large_patch16",
                    "c_radio_v2_vit_huge_patch16",
                    "c_radio_v3_vit_large_patch16_reg4_dinov2",
                    "c_radio_v3_vit_base_patch16_reg4_dinov2",
                    "c_radio_v3_vit_huge_patch16_reg4_dinov2",
                    "vit_l_14_siglip_clipa_224",
                    "vit_l_14_siglip_clipa_336",
                    "vit_h_14_siglip_clipa_224",
                    "mit_b0",
                    "mit_b1",
                    "mit_b2",
                    "mit_b3",
                    "mit_b4",
                    "mit_b5"
                  ],
                  "title": "Backbone architectures",
                  "type": "categorical"
                }
              },
              "type": "collection"
            },
            "head": {
              "automl_disabled_parameters": [
                "distill.teacher.head.custom_args",
                "distill.teacher.head.loss",
                "distill.teacher.head.topk"
              ],
              "automl_enabled": false,
              "default": {
                "binary": false,
                "in_channels": 448,
                "loss": {
                  "label_smooth_val": 0.0,
                  "type": "CrossEntropyLoss"
                },
                "topk": [
                  1
                ],
                "type": "TAOLinearClsHead"
              },
              "properties": {
                "binary": {
                  "default": false,
                  "description": "Flag to specify binary classification",
                  "type": "bool"
                },
                "custom_args": {
                  "automl_enabled": false,
                  "description": "custom head arguments",
                  "type": "collection"
                },
                "in_channels": {
                  "default": 448,
                  "description": "Number of backbone input channels to head",
                  "type": "int"
                },
                "loss": {
                  "automl_enabled": false,
                  "default": {
                    "label_smooth_val": 0.0,
                    "type": "CrossEntropyLoss"
                  },
                  "properties": {
                    "label_smooth_val": {
                      "default": 0.0,
                      "description": "Label smoothing value",
                      "maximum": 1.0,
                      "minimum": 0.0,
                      "type": "float"
                    },
                    "type": {
                      "default": "CrossEntropyLoss",
                      "description": "Loss type",
                      "enum": [
                        "CrossEntropyLoss"
                      ],
                      "type": "categorical"
                    }
                  },
                  "type": "collection"
                },
                "topk": {
                  "automl_enabled": false,
                  "default": [
                    1
                  ],
                  "description": "k value for Topk accuracy",
                  "type": "list"
                },
                "type": {
                  "default": "TAOLinearClsHead",
                  "description": "Type of classification head",
                  "enum": [
                    "TAOLinearClsHead",
                    "LogisticRegressionHead"
                  ],
                  "type": "categorical"
                }
              },
              "type": "collection"
            }
          },
          "title": "teacher",
          "type": "collection"
        },
        "use_mlp": {
          "default": true,
          "description": "Flag to use MLP for projection",
          "type": "bool"
        }
      },
      "type": "collection"
    },
    "encryption_key": {
      "default": "",
      "description": "Key for encrypting model checkpoints",
      "title": "Encryption key",
      "type": "string"
    },
    "model": {
      "automl_disabled_parameters": [
        "model.backbone",
        "model.head"
      ],
      "automl_enabled": false,
      "default": {
        "backbone": {
          "freeze_backbone": false,
          "freeze_norm": false,
          "pretrained_backbone_path": "",
          "type": "fan_small_12_p4_hybrid"
        },
        "head": {
          "binary": false,
          "in_channels": 448,
          "loss": {
            "label_smooth_val": 0.0,
            "type": "CrossEntropyLoss"
          },
          "topk": [
            1
          ],
          "type": "TAOLinearClsHead"
        }
      },
      "properties": {
        "backbone": {
          "automl_default_parameters": [
            "model.backbone.freeze_backbone",
            "model.backbone.freeze_norm"
          ],
          "automl_enabled": false,
          "default": {
            "freeze_backbone": false,
            "freeze_norm": false,
            "pretrained_backbone_path": "",
            "type": "fan_small_12_p4_hybrid"
          },
          "properties": {
            "freeze_backbone": {
              "automl_enabled": true,
              "default": false,
              "description": "Flag to freeze backbone",
              "type": "bool"
            },
            "freeze_norm": {
              "automl_enabled": true,
              "default": false,
              "description": "Flag to freeze norm",
              "type": "bool"
            },
            "pretrained_backbone_path": {
              "default": "",
              "description": "Path to the pretrained model",
              "type": "string"
            },
            "type": {
              "default": "fan_small_12_p4_hybrid",
              "description": "Backbone architure",
              "enum": [
                "faster_vit_0_224",
                "faster_vit_1_224",
                "faster_vit_2_224",
                "faster_vit_3_224",
                "faster_vit_4_224",
                "faster_vit_5_224",
                "faster_vit_6_224",
                "faster_vit_4_21k_224",
                "faster_vit_4_21k_384",
                "faster_vit_4_21k_512",
                "faster_vit_4_21k_768",
                "fan_tiny_12_p16_224",
                "fan_small_12_p16_224_se_attn",
                "fan_small_12_p16_224",
                "fan_base_18_p16_224",
                "fan_large_24_p16_224",
                "fan_tiny_8_p4_hybrid",
                "fan_small_12_p4_hybrid",
                "fan_base_16_p4_hybrid",
                "fan_large_16_p4_hybrid",
                "fan_xlarge_16_p4_hybrid",
                "fan_swin_tiny_patch4_window7_224",
                "fan_swin_small_patch4_window7_224",
                "fan_swin_base_patch4_window7_224",
                "fan_swin_large_patch4_window7_224",
                "vit_large_patch14_dinov2_swiglu",
                "vit_large_patch14_dinov2_swiglu_legacy",
                "vit_giant_patch14_reg4_dinov2_swiglu",
                "efficientvit_b0",
                "efficientvit_b1",
                "efficientvit_b2",
                "efficientvit_b3",
                "efficientvit_l0",
                "efficientvit_l1",
                "efficientvit_l2",
                "efficientvit_l3",
                "vit_base_patch16",
                "vit_large_patch16",
                "vit_huge_patch14",
                "convnext_tiny",
                "convnext_small",
                "convnext_base",
                "convnext_large",
                "convnext_xlarge",
                "convnextv2_atto",
                "convnextv2_femto",
                "convnextv2_pico",
                "convnextv2_nano",
                "convnextv2_tiny",
                "convnextv2_base",
                "convnextv2_large",
                "convnextv2_huge",
                "hiera_tiny_224",
                "hiera_small_224",
                "hiera_base_224",
                "hiera_base_plus_224",
                "hiera_large_224",
                "hiera_huge_224",
                "resnet_18",
                "resnet_34",
                "resnet_50",
                "resnet_101",
                "resnet_152",
                "resnet_18d",
                "resnet_34d",
                "resnet_50d",
                "resnet_101d",
                "resnet_152d",
                "swin_tiny_patch4_window7_224",
                "swin_small_patch4_window7_224",
                "swin_base_patch4_window7_224",
                "swin_large_patch4_window7_224",
                "swin_base_patch4_window12_384",
                "swin_large_patch4_window12_384",
                "gc_vit_xxtiny",
                "gc_vit_xtiny",
                "gc_vit_tiny",
                "gc_vit_small",
                "gc_vit_base",
                "gc_vit_large",
                "gc_vit_base_384",
                "gc_vit_large_384",
                "edgenext_xx_small",
                "edgenext_x_small",
                "edgenext_small",
                "edgenext_base",
                "edgenext_xx_small_bn_hs",
                "edgenext_x_small_bn_hs",
                "edgenext_small_bn_hs",
                "c_radio_p1_vit_huge_patch16_mlpnorm",
                "c_radio_p2_vit_huge_patch16_mlpnorm",
                "c_radio_p3_vit_huge_patch16_mlpnorm",
                "c_radio_v2_vit_base_patch16",
                "c_radio_v2_vit_large_patch16",
                "c_radio_v2_vit_huge_patch16",
                "c_radio_v3_vit_large_patch16_reg4_dinov2",
                "c_radio_v3_vit_base_patch16_reg4_dinov2",
                "c_radio_v3_vit_huge_patch16_reg4_dinov2",
                "vit_l_14_siglip_clipa_224",
                "vit_l_14_siglip_clipa_336",
                "vit_h_14_siglip_clipa_224",
                "mit_b0",
                "mit_b1",
                "mit_b2",
                "mit_b3",
                "mit_b4",
                "mit_b5"
              ],
              "title": "Backbone architectures",
              "type": "categorical"
            }
          },
          "type": "collection"
        },
        "head": {
          "automl_disabled_parameters": [
            "model.head.custom_args",
            "model.head.loss",
            "model.head.topk"
          ],
          "automl_enabled": false,
          "default": {
            "binary": false,
            "in_channels": 448,
            "loss": {
              "label_smooth_val": 0.0,
              "type": "CrossEntropyLoss"
            },
            "topk": [
              1
            ],
            "type": "TAOLinearClsHead"
          },
          "properties": {
            "binary": {
              "default": false,
              "description": "Flag to specify binary classification",
              "type": "bool"
            },
            "custom_args": {
              "automl_enabled": false,
              "description": "custom head arguments",
              "type": "collection"
            },
            "in_channels": {
              "default": 448,
              "description": "Number of backbone input channels to head",
              "type": "int"
            },
            "loss": {
              "automl_enabled": false,
              "default": {
                "label_smooth_val": 0.0,
                "type": "CrossEntropyLoss"
              },
              "properties": {
                "label_smooth_val": {
                  "default": 0.0,
                  "description": "Label smoothing value",
                  "maximum": 1.0,
                  "minimum": 0.0,
                  "type": "float"
                },
                "type": {
                  "default": "CrossEntropyLoss",
                  "description": "Loss type",
                  "enum": [
                    "CrossEntropyLoss"
                  ],
                  "type": "categorical"
                }
              },
              "type": "collection"
            },
            "topk": {
              "automl_enabled": false,
              "default": [
                1
              ],
              "description": "k value for Topk accuracy",
              "type": "list"
            },
            "type": {
              "default": "TAOLinearClsHead",
              "description": "Type of classification head",
              "enum": [
                "TAOLinearClsHead",
                "LogisticRegressionHead"
              ],
              "type": "categorical"
            }
          },
          "type": "collection"
        }
      },
      "type": "collection"
    },
    "model_name": {
      "default": "",
      "description": "Name of model if invoking task via :code:`model_agnostic`",
      "title": "Model name",
      "type": "string"
    },
    "quantize": {
      "automl_disabled_parameters": [
        "quantize.layers",
        "quantize.skip_names",
        "quantize.backend_kwargs"
      ],
      "automl_enabled": false,
      "default": {
        "algorithm": "minmax",
        "backend": "torchao",
        "backend_kwargs": {},
        "device": "cuda",
        "layers": [],
        "mode": "weight_only_ptq",
        "model_path": "",
        "results_dir": "",
        "skip_names": []
      },
      "properties": {
        "algorithm": {
          "default": "minmax",
          "description": "Calibration/optimization algorithm. Used by ModelOpt backends (modelopt.pytorch and modelopt.onnx). Ignored by torchao backend.",
          "enum": [
            "minmax",
            "max",
            "entropy",
            "awq_clip",
            "awq_lite",
            "awq_full",
            "rtn_dq"
          ],
          "title": "Calibration algorithm",
          "type": "categorical"
        },
        "backend": {
          "default": "torchao",
          "description": "The quantization backend to use",
          "enum": [
            "modelopt.pytorch",
            "torchao",
            "modelopt.onnx"
          ],
          "title": "Quantization backend",
          "type": "categorical"
        },
        "backend_kwargs": {
          "automl_enabled": false,
          "description": "Additional keyword arguments to pass to the backend",
          "title": "Backend kwargs",
          "type": "collection"
        },
        "device": {
          "default": "cuda",
          "description": "Device to use for calibration. Accepts 'cuda' (uses default GPU), 'cpu', 'trt' (TensorRT for ONNX backend), or specific GPU device like 'cuda:0', 'cuda:1', etc. If 'cuda' is specified but no GPU is available, the framework will automatically fall back to 'cpu'. Note: 'trt' is only supported by the modelopt.onnx backend for ONNX Runtime with TensorRT execution provider.",
          "pattern": "^(cuda|cpu|trt|cuda:[0-9]+)$",
          "title": "Calibration device",
          "type": "string"
        },
        "layers": {
          "automl_enabled": false,
          "default": [],
          "description": "List of per-module quantization configurations. Each entry specifies which modules to quantize and their data types. This is the primary way to configure quantization.",
          "title": "Layer quantization configs",
          "type": "list"
        },
        "mode": {
          "default": "weight_only_ptq",
          "description": "The quantization mode to use",
          "enum": [
            "static_ptq",
            "weight_only_ptq"
          ],
          "title": "Quantization mode",
          "type": "categorical"
        },
        "model_path": {
          "default": "",
          "description": "Path to the model to be quantized. For ONNX backend, path to ONNX file.",
          "title": "Model path",
          "type": "string"
        },
        "results_dir": {
          "default": "",
          "description": "Path to where all the assets generated from a task are stored.",
          "title": "Results directory",
          "type": "string"
        },
        "skip_names": {
          "automl_enabled": false,
          "default": [],
          "description": "List of module/layer names or patterns to exclude from quantization",
          "title": "Skip names",
          "type": "list"
        }
      },
      "required": [
        "model_path",
        "results_dir"
      ],
      "type": "collection"
    },
    "results_dir": {
      "default": "",
      "description": "\n        Path to where all the assets generated from a task are stored.\n        ",
      "title": "Results directory",
      "type": "string"
    },
    "train": {
      "automl_disabled_parameters": [
        "train.gpu_ids",
        "train.cudnn",
        "train.optim",
        "train.tensorboard"
      ],
      "automl_enabled": false,
      "default": {
        "checkpoint_interval": 1,
        "checkpoint_interval_unit": "epoch",
        "clip_grad_norm": 2.0,
        "cudnn": {
          "benchmark": false,
          "deterministic": true
        },
        "ema_decay": 0.998,
        "enable_ema": false,
        "gpu_ids": [
          0
        ],
        "num_epochs": 10,
        "num_gpus": 1,
        "num_nodes": 1,
        "optim": {
          "betas": [
            0.9,
            0.999
          ],
          "lr": 6e-05,
          "momentum": 0.9,
          "monitor_name": "val_loss",
          "optim": "adamw",
          "policy": "linear",
          "policy_params": {
            "gamma": 0.1,
            "step_size": 30
          },
          "skip_names": [],
          "warmup_epochs": 0,
          "weight_decay": 0.01
        },
        "precision": "fp32",
        "pretrained_model_path": "",
        "results_dir": "",
        "resume_training_checkpoint_path": "",
        "seed": 1234,
        "tensorboard": {
          "enabled": false,
          "infrequent_logging_frequency": 2
        },
        "validation_interval": 1
      },
      "popular": [
        "validation_interval",
        "checkpoint_interval",
        "num_gpus",
        "num_nodes",
        "num_epochs",
        "gpu_ids"
      ],
      "properties": {
        "checkpoint_interval": {
          "default": 1,
          "description": "The interval (in epochs) at which a checkpoint will be saved. Helps resume training.",
          "minimum": 1,
          "popular": true,
          "title": "Checkpoint interval",
          "type": "int"
        },
        "checkpoint_interval_unit": {
          "default": "epoch",
          "description": "The unit of the checkpoint interval.",
          "enum": [
            "epoch",
            "step"
          ],
          "title": "Checkpoint interval unit",
          "type": "categorical"
        },
        "clip_grad_norm": {
          "default": 2.0,
          "description": "Gradient Norm",
          "title": "Grad norm",
          "type": "float"
        },
        "cudnn": {
          "automl_enabled": false,
          "default": {
            "benchmark": false,
            "deterministic": true
          },
          "properties": {
            "benchmark": {
              "default": false,
              "type": "bool"
            },
            "deterministic": {
              "default": true,
              "type": "bool"
            }
          },
          "type": "collection"
        },
        "ema_decay": {
          "default": 0.998,
          "description": "EMA decay",
          "title": "EMA decay",
          "type": "float"
        },
        "enable_ema": {
          "default": false,
          "description": "Flag to enable EMA",
          "type": "bool"
        },
        "gpu_ids": {
          "automl_enabled": false,
          "default": [
            0
          ],
          "description": "\n        List of GPU IDs to run the training on. The length of this list\n        must be equal to the number of gpus in train.num_gpus.",
          "popular": true,
          "title": "GPU IDs",
          "type": "list"
        },
        "num_epochs": {
          "default": 10,
          "description": "Number of epochs to run the training.",
          "maximum": Infinity,
          "minimum": 1,
          "popular": true,
          "title": "Number of epochs",
          "type": "int"
        },
        "num_gpus": {
          "default": 1,
          "description": "The number of GPUs to run the train job.",
          "minimum": 1,
          "popular": true,
          "title": "Number of GPUs",
          "type": "int"
        },
        "num_nodes": {
          "default": 1,
          "description": "Number of nodes to run the training on. If > 1, then multi-node is enabled.",
          "minimum": 1,
          "popular": true,
          "title": "Number of nodes",
          "type": "int"
        },
        "optim": {
          "automl_default_parameters": [
            "train.optim.lr",
            "train.optim.momentum",
            "train.optim.weight_decay",
            "train.optim.betas"
          ],
          "automl_disabled_parameters": [
            "train.optim.policy_params",
            "train.optim.skip_names"
          ],
          "automl_enabled": false,
          "default": {
            "betas": [
              0.9,
              0.999
            ],
            "lr": 6e-05,
            "momentum": 0.9,
            "monitor_name": "val_loss",
            "optim": "adamw",
            "policy": "linear",
            "policy_params": {
              "gamma": 0.1,
              "step_size": 30
            },
            "skip_names": [],
            "warmup_epochs": 0,
            "weight_decay": 0.01
          },
          "properties": {
            "betas": {
              "automl_enabled": true,
              "default": [
                0.9,
                0.999
              ],
              "description": "coefficients used for computing running averages on adamw",
              "type": "list"
            },
            "lr": {
              "automl_enabled": true,
              "default": 6e-05,
              "description": "Optimizer learning rate",
              "maximum": Infinity,
              "minimum": 0.0,
              "type": "float"
            },
            "momentum": {
              "automl_enabled": true,
              "default": 0.9,
              "description": "The momentum for the AdamW optimizer.",
              "math_cond": "> 0.0",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "momentum - AdamW",
              "type": "float"
            },
            "monitor_name": {
              "default": "val_loss",
              "description": "Monitor Name",
              "type": "string"
            },
            "optim": {
              "default": "adamw",
              "description": "Optimizer",
              "enum": [
                "adamw",
                "adam",
                "sgd"
              ],
              "type": "categorical"
            },
            "policy": {
              "default": "linear",
              "description": "Optimizer policy",
              "enum": [
                "linear",
                "step",
                "cosine",
                "multistep"
              ],
              "type": "categorical"
            },
            "policy_params": {
              "automl_enabled": false,
              "default": {
                "gamma": 0.1,
                "step_size": 30
              },
              "description": "Optimizer policy parameters",
              "type": "collection"
            },
            "skip_names": {
              "automl_enabled": false,
              "default": [],
              "description": "layers names which do not need weight decay",
              "type": "list"
            },
            "warmup_epochs": {
              "default": 0,
              "description": "Warmup epochs.",
              "maximum": Infinity,
              "minimum": 0,
              "type": "int"
            },
            "weight_decay": {
              "automl_enabled": true,
              "default": 0.01,
              "description": "The weight decay coefficient.",
              "math_cond": "> 0.0",
              "maximum": 1.0,
              "minimum": 0.0,
              "title": "weight decay",
              "type": "float"
            }
          },
          "type": "collection"
        },
        "precision": {
          "default": "fp32",
          "description": "Precision",
          "enum": [
            "fp16",
            "fp32"
          ],
          "type": "categorical"
        },
        "pretrained_model_path": {
          "default": "",
          "description": "Pretrained model path",
          "title": "pretrained model path",
          "type": "string"
        },
        "results_dir": {
          "default": "",
          "description": "\n        Path to where all the assets generated from a task are stored.\n        ",
          "title": "Results directory",
          "type": "string"
        },
        "resume_training_checkpoint_path": {
          "default": "",
          "description": "Path to the checkpoint to resume training from.",
          "title": "Resume checkpoint path",
          "type": "string"
        },
        "seed": {
          "default": 1234,
          "description": "The seed for the initializer in PyTorch. If < 0, disable fixed seed.",
          "maximum": Infinity,
          "minimum": -1,
          "title": "Seed for randomization",
          "type": "int"
        },
        "tensorboard": {
          "automl_enabled": false,
          "default": {
            "enabled": false,
            "infrequent_logging_frequency": 2
          },
          "properties": {
            "enabled": {
              "default": false,
              "description": "Flag to enable tensorboard",
              "type": "bool"
            },
            "infrequent_logging_frequency": {
              "default": 2,
              "description": "infrequent_logging_frequency",
              "maximum": Infinity,
              "minimum": 0,
              "type": "int"
            }
          },
          "type": "collection"
        },
        "validation_interval": {
          "default": 1,
          "description": "\n        The interval (in epochs) at which a evaluation\n        will be triggered on the validation dataset.",
          "minimum": 1,
          "popular": true,
          "title": "Validation interval",
          "type": "int"
        }
      },
      "type": "collection"
    },
    "wandb": {
      "automl_disabled_parameters": [
        "wandb.tags"
      ],
      "automl_enabled": false,
      "default": {
        "enable": true,
        "entity": "",
        "group": "",
        "name": "TAO Toolkit Training",
        "project": "TAO Toolkit",
        "reinit": false,
        "run_id": "",
        "save_code": false,
        "sync_tensorboard": false,
        "tags": [
          "tao-toolkit"
        ]
      },
      "properties": {
        "enable": {
          "default": true,
          "type": "bool"
        },
        "entity": {
          "default": "",
          "type": "string"
        },
        "group": {
          "default": "",
          "type": "string"
        },
        "name": {
          "default": "TAO Toolkit Training",
          "type": "string"
        },
        "project": {
          "default": "TAO Toolkit",
          "type": "string"
        },
        "reinit": {
          "default": false,
          "type": "bool"
        },
        "run_id": {
          "default": "",
          "type": "string"
        },
        "save_code": {
          "default": false,
          "type": "bool"
        },
        "sync_tensorboard": {
          "default": false,
          "type": "bool"
        },
        "tags": {
          "automl_enabled": false,
          "default": [
            "tao-toolkit"
          ],
          "type": "list"
        }
      },
      "type": "collection"
    }
  },
  "required": [
    "quantize"
  ],
  "type": "object",
  "x_tao_schema": {
    "action": "train",
    "core_module": "classification_pyt",
    "model": "classification-pyt",
    "network_arch": "classification_pyt",
    "schema_action": "train",
    "schema_version": 1,
    "source": "tao-core dataclass config"
  }
}

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Declares allowed tools, Needs Python, Needs Node.js

The author says: Requires docker + nvidia-container-toolkit.

Security

What our scanner noticed· Medium risk

Large base64-encoded blobs

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Created by

NVIDIA Corporation

Publisher

NVIDIA (Official)

Source

NVIDIA/skills · skills/tao-train-image-classification · commit fce3cca

Licence

Apache License 2.0 (redistribution allowed)

LICENSE (added by AutoGineer from LICENSE-APACHE)
ATTRIBUTION.md (written by AutoGineer)

Synced

2026-11-09

AutoGineer is not affiliated with or endorsed by NVIDIA Corporation or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download tao-train-image-classification.zip