Image Classification PyT

Image Classification PyT is a PyTorch-based image-classification model included in TAO. It supports the following tasks:

• train

• evaluate

• inference

• export

These tasks can be invoked from the TAO Launcher using the following convention on the command-line:

            

            
tao model classification_pyt <sub_task> <args_per_subtask>
        

Where, args_per_subtask are the command-line arguments required for a given subtask. Each subtask is explained in detail in the following sections.

TAO has been updated to the latest MMPretrain version instead of the deprecated MMClassification. .. Note:: Image Classification (PyT) is based off of MMPretrain. Hence, most parameters are adopted from the MMPretrain 1.x format.

See the Data Annotation Format page for more information about the data format for image classification.

The train classification experiment specification consists of three main components:

• dataset

• train

• model

Here is an example of dataset specification file for classification PyT with a FAN backbone:

            

            
dataset:
  data:
    samples_per_gpu: 128
    workers_per_gpu: 8
    train:
      data_prefix: "/raid/ImageNet2012/ImageNet2012/train"
      pipeline: # Augmentations alone
        - type: RandomResizedCrop
          scale: 224
        - type: RandomFlip
          prob: 0.5
          direction: "horizontal"
        - type: ColorJitter
          brightness: 0.4
          contrast: 0.4
          saturation: 0.4
        - type: RandomErasing
          erase_prob: 0.3
    val:
      data_prefix: /raid/ImageNet2012/ImageNet2012/val
    test:
      data_prefix: /raid/ImageNet2012/ImageNet2012/val
        

The table below describes the configurable parameters in dataset.

data

pipeline

The following is an example pipeline config with different augmentations:

            

            
pipeline: # Augmentations alone
  - type: RandomResizedCrop
    scale: 224
    backend: "pillow"
  - type: RandomFlip
    prob: 0.5
    direction: "horizontal"
  - type: ColorJitter
    brightness: 0.4
    contrast: 0.4
    saturation: 0.4
  - type: RandomErasing
    erase_prob: 0.3
        

Some of the widely adopted augmentations and the parameters are listed below. More information, refer to the MMPretrain documentation for transforms

Here is an example of a train specification file for Image Classification PyT:

            

            
train:
  train_config:
    runner:
      max_epochs: 300
    checkpoint_config:
      interval: 1
    logging:
      interval: 5000
    validate: True
    evaluation:
      interval: 10
    custom_hooks:
      - type: "EMAHook"
        momentum: 0.00008
        priority: "ABOVE_NORMAL"
    lr_config:
      policy: CosineAnnealing
      T_max: 95
      by_epoch: True
      begin: 5
      end: 100
    optimizer:
        type: AdamW
        lr: 0.005
        weight_decay: 0.05
        

The table below describes the configurable parameters in the train specification.

train_config

optimizer

The following optimizers are supported:

SGD

AdamW

lr_config

The lr_config parameter defines the parameters for the learning-rate scheduler. Some of the widely adopted learning-rate schedulers and the parameters are listed below. More information, refer to the MMPretrain documentation for lr schedulers

CosineAnnealingLR

MultiStepLR

LinearLR

PolyLR

StepLR

custom_hooks

The following is an example of how a custom hook from the MMPretrain to Hydra config is provided for EMAHook`:

• MMPretrain config:

  Copy

  Copied!

              
  
              
  custom_hooks = [
    dict(type='EMAHook', interval=100, priority='HIGH')]
          

• Equivalent TAO Hydra config:

  Copy

  Copied!

              
  
              
  custom_hooks:
    - type: "EMAHook"
      momentum: 0.00008
      priority: "ABOVE_NORMAL"
          

For more detail on custom_hooks, refer to the MMPretrain documentation for custom hooks.

Here is an example model-specification file for Image Classification PyT with a FAN backbone:

            

            
model:
  backbone:
    type: "fan_tiny_8_p4_hybrid"
    custom_args:
      drop_path_rate: 0.1
    freeze: False
    pretrained: <Path to pretrained weights>
  head:
    type: "TAOLinearClsHead"
    num_classes: 1000
    custom_args:
      head_init_scale: 1
    loss:
      type: LabelSmoothLoss
      label_smooth_val: 0.1
      mode: 'original'
  train_cfg:
    augments:
      - type: Mixup
        alpha: 0.8
      - type: CutMix
        alpha: 1.0
        

The model parameter primarily configures the backbone and head.

Foundation Models

            

            
model:
  backbone:
    type: "ViT-B-32"
    custom_args:
      drop_path_rate: 0.1
    freeze: False
    pretrained: laion400m_e31
  head:
    type: LinearClsHead
    num_classes: 1000
    in_channels: 512
    loss:
      type: CrossEntropyLoss
      loss_weight: 1.0
      use_soft: False
    topk: [1, 5]
        

Subset of the supported arch and the pre-train datasets. Please note that the in_channels should be updated under the head :

• CLIP Image Backbones:

• EVA - CLIP Image Backbones:

head

lr_head

Logistic Regression head is defined by the following parameters:

Logistic Regression head involves freezing the backbone weights. Please note that the freeze parameter should be updated to True under the backbone config as under:

            

            
model:
  backbone:
    type: "ViT-B-32"
    custom_args:
      drop_path_rate: 0.1
    freeze: True
    pretrained: laion400m_e31
  head:
    type: "LogisticRegressionHead"
    num_classes: 1000
    lr_head:
      C: 0.316
      max_iter: 5000
        

train_cfg

Mixup

CutMix

loss

Some Important Losses for classification losses are shown below. Please note that all supported losses in MMPretrain can be used by following the Hydra config for TAO. For a list of MMPretrain losses, refer to the losses_mmpretrain documentation.

LabelSmoothLoss

CrossEntropyLoss

Use the tao model classification_pyt train command to train a classification pytorch model:

            

            
tao model classification_pyt train [-h] -e <experiment_spec_file>
                             [results_dir=<global_results_dir>]
                             [model.<model_option>=<model_option_value>]
                             [dataset.<dataset_option>=<dataset_option_value>]
                             [train.<train_option>=<train_option_value>]
                             [train.gpu_ids=<gpu indices>]
                             [train.num_gpus=<number of gpus>]
        

Required Arguments

The only required argument is the path to the experiment spec:

• -e, --experiment_spec: The experiment specification file to set up the training experiment

Optional Arguments

You can set optional arguments to override the option values in the experiment spec file.

• -h, --help: Show this help message and exit.

• model.<model_option>: The model options.

• dataset.<dataset_option>: The dataset options.

• train.<train_option>: The train options.

• train.train_config.optimizer.<optim_option>: The optimizer options

The evaluate config defines the hyperparameters of the evaluation process. The following is an example config:

            

            
evaluate:
  checkpoint: /path/to/model.pth
  topk: 1
        

After the model has been trained using the experiment config file and by following the steps to train a model, the next step is to evaluate this model on a test set to measure the accuracy of the model. TAO includes the tao model classification_pyt evaluate command to do this.

The classification app computes evaluation loss and Top-k accuracy.

After training, the model is stored in the output directory of your choice in results_dir.

            

            
evaluate:
  checkpoint: /path/to/model.pth
        

            

            
tao model classification_pyt evaluate [-h] -e <experiment_spec>
                             evaluate.checkpoint=<model to be evaluated>
                             results_dir=<path to results dir>
                             [evaluate.<evaluate_option>=<evaluate_option_value>]
                             [evaluate.gpu_ids=<gpu indices>]
                             [evaluate.num_gpus=<number of gpus>]
        

Required Arguments

• -e, --experiment_spec: The experiment spec file to set up the evaluation experiment

• evaluate.checkpoint: The .pth model to be evaluated.

• results_dir: The path where the results will be stored

Optional Arguments

• evaluate.<evaluate_option>: The evaluate options.

For classification, tao model classification_pyt inference saves a .csv file containing the image paths and the corresponding labels for multiple images. TensorRT Python inference can also be enabled.

            

            
inference:
  checkpoint: /path/to/model.pth
        

            

            
tao model classification_pyt inference [-h] -e <experiment_spec_file>
                             inference.checkpoint=<model to be inferenced>
                             results_dir=<path to results dir>
                             [inference.<inference_option>=<inference_option_value>]
                             [inference.gpu_ids=<gpu indices>]
                             [inference.num_gpus=<number of gpus>]
        

Required Arguments

• -e, --experiment_spec: The experiment spec file to set up the inference experiment

• inference.checkpoint: The .pth model to inference.

• results_dir: The path where the results will be stored

Optional Arguments

• inference.<inference_option>: The inference options.

Exporting the model decouples the training process from inference and allows conversion to TensorRT engines outside the TAO environment. TensorRT engines are specific to each hardware configuration and should be generated for each unique inference environment. The exported model may be used universally across training and deployment hardware. The exported model format is referred to as .onnx.

The export parameter defines the hyperparameters of the export process.

            

            
export:
  checkpoint: /path/to/model.pth
  onnx_file: /path/to/model.onnx
  opset_version: 12
  verify: False
  input_channel: 3
  input_width: 224
  input_height: 224
        

Here’s an example of the tao classification_pyt export command:

            

            
tao model classification_pyt export [-h] -e <experiment spec file>
                             export.checkpoint=<model to export>
                             export.onnx_file=<onnx path>
                             [export.<export_option>=<export_option_value>]
        

Required Arguments

• -e, --experiment_spec: The path to an experiment spec file

• export.checkpoint: The .pth model to export.

• export.onnx_file: The path where the .etlt or .onnx model is saved.

Optional Arguments

• export.<export_option>: The export options.

For TensorRT engine generation, validation, and INT8 calibration, refer to the TAO Deploy documentation.

Refer to the Integrating a Classification (TF1/TF2/PyTorch) Model page for more information about deploying a classification model with DeepStream.