Image Classification PyT

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

• train

• evaluate

• inference

• export

These tasks can be invoked from the TAO Toolkit 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 Toolkit 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
          size: 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
    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: BatchMixup
        alpha: 0.8
        num_classes: 1000
        prob: 0.5
      - type: BatchCutMix
        alpha: 1.0
        num_classes: 1000
        prob: 0.5
        

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

BatchMixup

BatchCutMix

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 Toolkit. 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 <spec file>
                              -r <result directory>
                              [-g <num GPUs>]
        

Required Arguments

• -r, --results_dir: The path to a folder where the experiment outputs should be written

• -e, --experiment_spec_file: The path to the experiment spec file

Optional Arguments

• -g, --gpus: The nubmer of GPUs to use for training. The default value is 1.

• -h, --help: Print the help message.

Sample Usage

Here’s an example of using the tao model classification_pyt train command:

            

            
tao model classification_pyt train -e /workspace/cats_dogs/spec/train_cats_dogs.yaml -r /workspace/output
        

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. The TAO toolkit 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_file>
                                      evaluate.checkpoint=<model to be evaluated>
                                      results_dir=<path to results dir>
                                      [-g <num gpus>]
        

Required Arguments

• -e, --experiment_spec_file: The path to the experiment spec file

Optional Arguments

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

• -g, --gpus: The number of GPUs for conducting evaluation

If you followed the example in training a classification model, run the evaluation:

            

            
tao model classification_pyt evaluate -e /path/to/classification_eval.yaml
        

TAO will evaluate for classification and produces the Top-K accuracy metric.

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>
                                      [-g <num gpus>]
        

Required Arguments

• -e, --experiment_spec_file: The path to the experiment spec file

Optional Arguments

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

• -g, --gpus: The number of GPUs to conduct the evaluation

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>
                  [-r <results_dir>]
                  export.checkpoint=<model to export>
                  export.onnx_file=<onnx path>
        

Required Arguments

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

Optional Arguments

• -r, --results_dir: The directory where the inference result is stored

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

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

Sample Usage

The following is a sample export command.

            

            
tao model classification_pyt export -e /path/to/spec.yaml export.checkpoint=/path/to/model.pth export.onnx_file=/path/to/model.onnx
        

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.