Tao Train Image Classification

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

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.