Tao Port Huggingface Model

# tao-pytorch commands
<model_name> <subtask> -e <experiment_spec.yaml> [hydra_overrides...]

# Examples:
classification_pyt train -e experiment_spec.yaml train.num_epochs=50 train.optim.lr=0.001
segformer evaluate -e experiment_spec.yaml evaluate.checkpoint=/path/to/model.pth
dino export -e experiment_spec.yaml export.onnx_file=/results/model.onnx

# tao-deploy commands (identical pattern)
classification_pyt gen_trt_engine -e gen_trt_engine.yaml gen_trt_engine.tensorrt.data_type=FP16
classification_pyt inference -e inference.yaml inference.trt_engine=/path/to/model.engine
classification_pyt evaluate -e evaluate.yaml evaluate.trt_engine=/path/to/model.engine

console_script main()
  → get_subtasks(scripts_package)         # discovers train.py, evaluate.py, etc. via pkgutil
  → command_line_parser(parser, subtasks)  # extracts: subtask, -e spec_file, unknown_args
  → launch(args, unknown_args, subtasks)   # orchestrates execution

# tao-pytorch entrypoint: nvidia_tao_pytorch/cv/<model_name>/entrypoint/<model_name>.py
from nvidia_tao_pytorch.cv.<model_name> import scripts
from nvidia_tao_pytorch.core.entrypoint import get_subtasks, command_line_parser, launch

def main():
    subtasks = get_subtasks(scripts)
    args, unknown_args = command_line_parser(subtasks)
    launch(vars(args), unknown_args, subtasks, network="<model_name>")

if __name__ == "__main__":
    main()

# tao-deploy entrypoint: nvidia_tao_deploy/cv/<model_name>/entrypoint/<model_name>.py
from nvidia_tao_deploy.cv.<model_name> import scripts
from nvidia_tao_deploy.cv.common.entrypoint.entrypoint_hydra import (
    get_subtasks, command_line_parser, launch
)

def main():
    subtasks = get_subtasks(scripts)
    args, unknown_args = command_line_parser(subtasks)
    launch(vars(args), unknown_args, subtasks, network="<model_name>")

if __name__ == "__main__":
    main()

1. ExperimentConfig dataclass defaults (schema)
2. YAML experiment spec values
3. CLI overrides (dot notation)

@hydra_runner(
    config_path=os.path.join(spec_root, "experiment_specs"),  # directory containing YAML
    config_name="experiment_spec",                              # YAML filename (no .yaml)
    schema=ExperimentConfig                                     # dataclass from tao-core
)
@monitor_status(name="<ModelName>", mode="train")
def main(cfg: ExperimentConfig) -> None:
    ...

@dataclass
class ExperimentConfig(CommonExperimentConfig):
    """Top-level config — drives both tao-pytorch and tao-deploy scripts."""
    model: ModelConfig              = DATACLASS_FIELD(ModelConfig())
    dataset: DatasetConfig          = DATACLASS_FIELD(DatasetConfig())
    train: TrainExpConfig           = DATACLASS_FIELD(TrainExpConfig())
    evaluate: EvalExpConfig         = DATACLASS_FIELD(EvalExpConfig())
    inference: InferenceExpConfig   = DATACLASS_FIELD(InferenceExpConfig())
    export: ExportExpConfig         = DATACLASS_FIELD(ExportExpConfig())
    gen_trt_engine: GenTrtEngineExpConfig = DATACLASS_FIELD(GenTrtEngineExpConfig())
    quantize: ModelQuantizationConfig = DATACLASS_FIELD(ModelQuantizationConfig())
    # Optional (task-specific):
    distill: DistillConfig          = DATACLASS_FIELD(DistillConfig())

model_name: Optional[str]     # for model-agnostic invocation
encryption_key: Optional[str] # checkpoint encryption
results_dir: Optional[str]    # top-level output directory
wandb: WandBConfig            # experiment tracking (enable, project, entity, tags)

num_gpus: int          # default=1
gpu_ids: List[int]     # default=[0]
num_nodes: int         # default=1
seed: int              # default=1234
num_epochs: int        # default=10
checkpoint_interval: int  # default=1
validation_interval: int  # default=1
resume_training_checkpoint_path: Optional[str]
results_dir: Optional[str]
cudnn:
    benchmark: bool    # default=False
    deterministic: bool # default=True

num_gpus: int
gpu_ids: List[int]
checkpoint: str        # MISSING — required
trt_engine: Optional[str]
results_dir: Optional[str]
batch_size: int        # default=-1 (auto)

results_dir: Optional[str]
gpu_id: int            # default=0 (singular — export is single-GPU)
checkpoint: str        # MISSING — required
onnx_file: str         # MISSING — output path
input_channel: int     # default=3
input_width: int       # default=960
input_height: int      # default=544
opset_version: int     # default=17
batch_size: int        # default=-1 (dynamic)

results_dir: Optional[str]
gpu_id: int            # default=0
onnx_file: str         # MISSING — input ONNX path
trt_engine: Optional[str] # output engine path
tensorrt:
    data_type: str     # FP32, FP16, or INT8
    workspace_size: int # default=1024 (MB)
    min_batch_size: int # default=1
    opt_batch_size: int # default=1
    max_batch_size: int # default=1
    calibration:       # for INT8 only
        cal_image_dir: List[str]
        cal_cache_file: str
        cal_batch_size: int
        cal_batches: int

encryption_key: tlt_encode
results_dir: ???  # User must provide

model:
  backbone:
    type: "vit_large_patch14_dinov2_swiglu"
    pretrained_backbone_path: null
    freeze_backbone: False
    freeze_norm: False
  head:
    type: TAOLinearClsHead
    in_channels: 1024   # Must match backbone output dim from model_params_mapping.py
    topk: [1, 5]
    loss:
      type: CrossEntropyLoss
      label_smooth_val: 0.0

dataset:
  dataset: "CLDataset"
  root_dir: ???  # Location of classes.txt
  num_classes: 1000
  img_size: 224
  batch_size: 128
  workers: 8
  shuffle: True
  augmentation:
    mean: [0.485, 0.456, 0.406]
    std: [0.229, 0.224, 0.225]
    random_flip:
      enable: True
      hflip_probability: 0.5
      vflip_probability: 0.0
    random_rotate:
      enable: False
    random_color:
      enable: False
    random_erase:
      enable: False
  train_dataset:
    images_dir: ${dataset.root_dir}/train
  val_dataset:
    images_dir: ${dataset.root_dir}/val
  test_dataset:
    images_dir: ${dataset.root_dir}/test

train:
  seed: 1234
  num_epochs: 25
  num_gpus: 1
  gpu_ids: [0]
  num_nodes: 1
  checkpoint_interval: 10
  validation_interval: 1
  resume_training_checkpoint_path: null
  clip_grad_norm: 2.0
  precision: fp32
  enable_ema: False
  optim:
    optim: adamw
    lr: 0.00006
    weight_decay: 0.05
    policy: cosine
    warmup_epochs: 5
    momentum: 0.9
  tensorboard:
    enabled: True

evaluate:
  checkpoint: ${results_dir}/train/<model_name>_model_latest.pth

inference:
  checkpoint: ${results_dir}/train/<model_name>_model_latest.pth

export:
  results_dir: ${results_dir}/export
  gpu_id: 0
  checkpoint: ${results_dir}/train/<model_name>_model_latest.pth
  onnx_file: ${export.results_dir}/<model_name>.onnx
  input_width: 224
  input_height: 224
  batch_size: -1
  opset_version: 17

gen_trt_engine:
  onnx_file: ${export.results_dir}/<model_name>.onnx
  trt_engine: ${results_dir}/trt/<model_name>.engine
  tensorrt:
    data_type: FP16
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 4
    max_batch_size: 8

dataset:
  train_data_sources:
    - image_dir: /data/train/images
      json_file: /data/train/annotations.json
  val_data_sources:
    - image_dir: /data/val/images
      json_file: /data/val/annotations.json
  test_data_sources:
    image_dir: /data/test/images
    json_file: /data/test/annotations.json
  infer_data_sources:
    image_dir: [/data/infer/images]
    classmap: /data/classmap.txt
  num_classes: 91
  augmentation:
    scales: [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
    input_mean: [0.485, 0.456, 0.406]
    input_std: [0.229, 0.224, 0.225]

model:
  backbone: "fan_small_12_p4_hybrid"
  num_queries: 300
  num_feature_levels: 4
  enc_layers: 6
  dec_layers: 6
  hidden_dim: 256

dataset:
  segment:
    palette:
      - label_id: 0
        rgb: [0, 0, 0]
        mapping_class: "background"
        seg_class: "background"
      - label_id: 1
        rgb: [128, 0, 0]
        mapping_class: "person"
        seg_class: "person"
  augmentation:
    mean: [0.485, 0.456, 0.406]
    std: [0.229, 0.224, 0.225]

model:
  backbone:
    type: "mit_b5"
  decode_head:
    in_channels: [64, 128, 320, 512]
    in_index: [0, 1, 2, 3]
    feature_strides: [4, 8, 16, 32]

results_dir: ???
gen_trt_engine:
  onnx_file: ???
  trt_engine: ???
  tensorrt:
    data_type: FP16
    workspace_size: 1024
    min_batch_size: 1
    opt_batch_size: 4
    max_batch_size: 8

results_dir: ???
inference:
  trt_engine: ???
  batch_size: 8
dataset:
  root_dir: ???         # For classes.txt lookup
  test_dataset:
    images_dir: ???
  augmentation:
    mean: [0.485, 0.456, 0.406]
    std: [0.229, 0.224, 0.225]

results_dir: ???
evaluate:
  trt_engine: ???
  batch_size: 8
model:
  head:
    topk: [1]
dataset:
  root_dir: ???
  test_dataset:
    images_dir: ???
  augmentation:
    mean: [0.485, 0.456, 0.406]
    std: [0.229, 0.224, 0.225]

results_dir/                          ← set by user
├── train/                            ← train script writes here
│   ├── lightning_logs/version_1/
│   │   ├── hparams.yaml
│   │   ├── metrics.csv
│   │   └── events.out.tfevents.*     ← TensorBoard data
│   ├── model_001.pth                 ← checkpoint at epoch 1
│   ├── model_010.pth                 ← checkpoint at epoch 10
│   ├── <model_name>_model_latest.pth ← symlink to latest (IMPORTANT)
│   └── status.json                   ← TAO status logger
│
├── evaluate/                         ← evaluate script writes here
│   ├── result.csv                    ← per-image predictions + GT
│   └── results.json                  ← aggregate metrics
│
├── inference/                        ← inference script writes here
│   └── result.csv                    ← per-image predictions
│
├── export/                           ← export script writes here
│   ├── <model_name>.onnx            ← ONNX model
│   ├── labels.txt                    ← class labels (optional)
│   └── nvdsinfer_config.yaml         ← DeepStream config (optional)
│
├── trt/                              ← gen_trt_engine writes here
│   └── <model_name>.engine           ← TensorRT engine
│
├── trt_infer/                        ← TRT inference writes here
│   └── result.csv
│
└── trt_eval/                         ← TRT evaluation writes here
    └── results.json

train → export:
  export.checkpoint = ${results_dir}/train/<model_name>_model_latest.pth

export → gen_trt_engine:
  gen_trt_engine.onnx_file = ${results_dir}/export/<model_name>.onnx

gen_trt_engine → inference:
  inference.trt_engine = ${results_dir}/trt/<model_name>.engine

gen_trt_engine → evaluate:
  evaluate.trt_engine = ${results_dir}/trt/<model_name>.engine

class MyModelPlModel(TAOLightningModule):
    def __init__(self, experiment_spec):
        super().__init__(experiment_spec)
        self.checkpoint_filename = "<model_name>_model"  # e.g., "classifier_model"

ModelCheckpoint(
    filename="model_{epoch:03d}",    # per-epoch naming
    every_n_epochs=checkpoint_interval,
    save_last="link",                # creates _latest symlink
    save_top_k=-1,                   # keep all
    save_on_train_epoch_end=True,
    dirpath=results_dir,
)

root_dir/
├── classes.txt          ← one class name per line, sorted alphabetically
├── train/
│   ├── class_a/
│   │   ├── img001.jpg
│   │   └── img002.jpg
│   └── class_b/
│       └── img003.jpg
├── val/
│   └── ...              ← same structure as train/
└── test/
    └── ...              ← same structure as train/

data_dir/
├── train/
│   ├── images/
│   └── annotations.json   ← COCO JSON format
├── val/
│   ├── images/
│   └── annotations.json
└── classmap.txt            ← class_name per line for inference

data_dir/
├── train/
│   ├── images/
│   └── masks/             ← PNG masks with label IDs as pixel values
├── val/
│   ├── images/
│   └── masks/
└── test/
    ├── images/
    └── masks/

mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]

User YAML spec                    Entrypoint launch()              Training script
─────────────                     ──────────────────               ───────────────
train:                       →    Reads num_gpus, gpu_ids     →   TAO_VISIBLE_DEVICES env var
  num_gpus: 4                     Sets TAO_VISIBLE_DEVICES         parsed by initialize_train_experiment()
  gpu_ids: [0,1,2,3]             Wraps with torchrun:              → trainer_kwargs['devices']
  num_nodes: 1                    torchrun --nproc-per-node=4       → Trainer(devices=[0,1,2,3],
                                                                              strategy='ddp_find_unused_parameters_true',
                                                                              sync_batchnorm=True)

trainer_kwargs = {
    'logger': [TensorBoardLogger(save_dir=results_dir, version=1, name="lightning_logs")],
    'devices': [0, 1, 2, 3],    # from TAO_VISIBLE_DEVICES
    'max_epochs': cfg.train.num_epochs,
    'check_val_every_n_epoch': cfg.train.validation_interval,
    'default_root_dir': results_dir,
    'accelerator': 'gpu',
    'enable_checkpointing': False,  # PLModel defines own ModelCheckpoint in configure_callbacks()
}

# evaluate.py
model_path, trainer_kwargs = initialize_evaluation_experiment(cfg, key)
dm = <DataModule>(cfg.dataset)
dm.setup(stage="test")
model = <ModelPlModel>.load_from_checkpoint(model_path, map_location="cpu", experiment_spec=cfg)
trainer = Trainer(**trainer_kwargs)
trainer.test(model, datamodule=dm)

# inference.py
model_path, trainer_kwargs = initialize_inference_experiment(cfg, key)
dm = <DataModule>(cfg.dataset)
dm.setup(stage="predict")
model = <ModelPlModel>.load_from_checkpoint(model_path, map_location="cpu", experiment_spec=cfg)
trainer = Trainer(**trainer_kwargs)
trainer.predict(model, datamodule=dm)

# Load model
sf_model = <ModelPlModel>.load_from_checkpoint(model_path, map_location="cpu", experiment_spec=cfg)
model = sf_model.model  # Extract the raw nn.Module (not the PLModel wrapper)
model.eval()
model.cuda()

# Create dummy input matching export config
dummy_input = torch.ones(batch_size, input_channel, input_height, input_width, device='cuda')

# Export via ONNXExporter
from nvidia_tao_pytorch.core.exporters import ONNXExporter
onnx_exporter = ONNXExporter()
onnx_exporter.export_model(
    model, batch_size, output_file, dummy_input,
    input_names=['input'], output_names=['output'],
    opset_version=cfg.export.opset_version,
    do_constant_folding=True
)

# Decrypt ONNX if encrypted
tmp_onnx_file, file_format = decode_model(cfg.gen_trt_engine.onnx_file)

# Initialize builder kwargs
engine_builder_kwargs, create_engine_kwargs = initialize_gen_trt_engine_experiment(cfg)

# Detect QDQ quantization
strongly_typed = is_qdq_quantized_onnx(tmp_onnx_file) if file_format == "onnx" else False

# Build engine
builder = <ModelName>EngineBuilder(
    **engine_builder_kwargs,
    workspace=cfg.gen_trt_engine.tensorrt.workspace_size,
    is_qat=False,
    strongly_typed=strongly_typed,
    data_format="channels_first",
    preprocess_mode="torch"        # Must match training normalization
)
builder.create_network(tmp_onnx_file, file_format)
builder.create_engine(**create_engine_kwargs)

# Load class mapping
classmap = os.path.join(cfg.dataset.root_dir, 'classes.txt')
mapping_dict = {line.rstrip(): idx for idx, line in enumerate(sorted(open(classmap)))}

# Create TRT inferencer
trt_infer = ClassificationInferencer(
    cfg.inference.trt_engine,
    data_format="channel_first",
    batch_size=cfg.inference.batch_size
)

# Create NumPy dataloader
dl = ClassificationLoader(
    input_shape=trt_infer.input_tensors[0].shape,  # From TRT engine
    data_paths=[cfg.dataset.test_dataset.images_dir],
    class_mapping=mapping_dict,
    is_inference=True,
    batch_size=cfg.inference.batch_size,
    image_mean=cfg.dataset.augmentation.mean,       # Must match training
    image_std=cfg.dataset.augmentation.std           # Must match training
)

# Run inference and write results
with open(f"{cfg.results_dir}/result.csv", 'w') as csv_f:
    for imgs, _ in dl:
        y_pred = trt_infer.infer(imgs)
        class_indices = np.argmax(y_pred, axis=1)
        # Write to CSV

User sets: encryption_key: "tlt_encode" in YAML
  → initialize_train_experiment() calls TLTPyTorchCookbook.set_passphrase(key)
  → ModelCheckpoint saves .pth files (unencrypted) or .tlt files (encrypted)
  → Export can produce .etlt (encrypted ONNX) if key is set
  → Deploy decode_model() decrypts .etlt back to ONNX before TRT build