Image Classification (TF2)

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

Below is a sample for the classification spec file. It has seven major components–dataset, model, train, evaluate, inference, prune, and export–as well as parameters for the encryption key (encryption_key) and results directory (results_dir).

            

            
results_dir: '/workspace/results_dir'
encryption_key: 'nvidia_tlt'
dataset:
  train_dataset_path: "/workspace/tao-experiments/data/split/train"
  val_dataset_path: "/workspace/tao-experiments/data/split/val"
  preprocess_mode: 'torch'
  num_classes: 20
  augmentation:
    enable_color_augmentation: True
    enable_center_crop: True
train:
  qat: False
  checkpoint_interval: 1
  checkpoint: '/path/to/model.tlt'
  batch_size_per_gpu: 64
  num_epochs: 80
  optim_config:
    optimizer: 'sgd'
  lr_config:
    scheduler: 'cosine'
    learning_rate: 0.05
    soft_start: 0.05
  reg_config:
    type: 'L2'
    scope: ['conv2d', 'dense']
    weight_decay: 0.00005
model:
  backbone: 'efficientnet-b0'
  input_width: 224
  input_height: 224
  input_channels: 3
  input_image_depth: 8
evaluate:
  dataset_path: '/workspace/tao-experiments/data/split/test'
  checkpoint: 'EVALMODEL'
  top_k: 3
  batch_size: 256
  n_workers: 8
inference:
  checkpoint: 'EVALMODEL'
  image_dir: '/workspace/tao-experiments/data/split/test/aeroplane'
  classmap: 'RESULTSDIR/classmap.json'
export:
  checkpoint: 'EVALMODEL'
  onnx_file: 'EXPORTDIR/efficientnet-b0.onnx'
  data_type: 'fp32'
        

The format of the spec file is YAML. The top level structure of the spec file is summarized in the table below:

Model Config

The table below describes the parameters in the model config.

Dataset Config

The table below describes the parameters in the data config.

The table below describes the parameters in the dataset.augmentation config.

Inference Config

The table below defines the configurable parameters for evaluating a classification model.

Evaluation Config

The table below defines the configurable parameters for evaluating a classification model.

Training Config

This section defines the configurable parameters for the classification model trainer.

Learning Rate Scheduler

The parameter lr_config defines the parameters for learning rate scheduler The learning rate scheduler can be either step, soft_anneal or cosine.

Step

The parameter step defines the step learning rate scheduler.

Soft Annealing

The parameter soft_anneal defines the soft annealing learning rate scheduler.

Cosine

The parameter cosine defines the cosine learning rate scheduler.

Optimizer Config

Three types of optimizers are supported: Adam, SGD and RMSProp. Only one type should be specified in the spec file using optimizer

The Adam optimizer parameters are summarized in the table below.

The SGD optimizer parameters are summarized in the table below.

The RMSProp optimizer parameters are summarized in the table below.

BatchNorm Config

Pruning Config

The prune configuration defines the pruning process for a trained model. A detailed description is summarized in the table below.

Export Config

The export configuration contains the parameters for exporting a .tlt model to an .onnx model, which can be used for deployment.

Use the tao model classification_tf2 train command to tune a pre-trained model:

            

            
tao model classification_tf2 train [-h] -e <spec file>
                              [--gpus <num GPUs>]
                              [--gpu_index <gpu_index>]
                              [--log_file <log_file_path>]
        

Required Arguments

• -e, --experiment_spec_file: Path to the experiment spec file.

Optional Arguments

• --gpus: Number of GPUs to use and processes to launch for training. The default value is 1.

• --gpu_index: The GPU indices used to run the training. We can specify the GPU indices used to run training when the machine has multiple GPUs installed.

• --log_file: Path to the log file. Defaults to stdout.

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

Input Requirement

• Input size: 3 * H * W (W, H >= 32)

• Input format: JPG, JPEG, PNG

Sample Usage

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

            

            
tao model classification_tf2 train -e /workspace/spec.yaml --gpus 2
        

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 classification app computes evaluation loss, Top-k accuracy, precision, and recall as metrics. Evaluate a model using the tao model classification_tf2 evaluate command:

            

            
tao model classification_tf2 evaluate [-h] -e <experiment_spec_file>
                                 [--gpu_index <gpu_index>]
                                 [--log_file <log_file>]
        

Required Arguments

• -e, --experiment_spec_file: Path to the experiment spec file.

Optional Arguments

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

• --gpu_index: The GPU indices used to run the training. We can specify the GPU indices used to run training when the machine has multiple GPUs installed.

• --log_file: Path to the log file. Defaults to stdout.

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

            

            
tao model classification_tf2 evaluate -e classification_spec.yaml
        

TAO evaluates for classification and produces the following metrics:

• Loss

• Top-K accuracy

• Precision (P): TP / (TP + FP)

• Recall (R): TP / (TP + FN)

• Confusion Matrix

The tao model classification_tf2 inference command runs the inference on a specified set of input images. For classification, tao model classification_tf2 inference provides class label output over the command-line for a single image or a csv file containing the image path and the corresponding labels for multiple images. TensorRT Python inference can also be enabled.

Execute tao model classification_tf2 inference on a classification model trained on TAO Toolkit.

            

            
tao model classification_tf2 inference [-h] -e <experiment_spec_file>
                                  [--gpu_index <gpu_index>]
                                  [--log_file <log_file>]
        

Here are the arguments of the tao model classification_tf2 inference tool:

Required arguments

• -e, --experiment_spec_file: Path to the experiment spec file.

Optional arguments

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

• --gpu_index: The GPU indices used to run the training. We can specify the GPU indices used to run training when the machine has multiple GPUs installed.

• --log_file: Path to the log file. Defaults to stdout.

Pruning removes parameters from the model to reduce the model size without compromising the integrity of the model itself using the tao model classification_tf2 prune command.

The tao model classification_tf2 prune command includes these parameters:

            

            
tao model classification_tf2 prune [-h] -e <experiment_spec_file>
        

Required Arguments

• -e, --experiment_spec_file: Path to the experiment spec file.

Optional Arguments

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

• --gpu_index: The GPU indices used to run the training. We can specify the GPU indices used to run training when the machine has multiple GPUs installed.

• --log_file: Path to the log file. Defaults to stdout.

After pruning, the model needs to be retrained. See Re-training the Pruned Model for more details.

Using the Prune Command

Here’s an example of using the tao model classification_tf2 prune command:

            

            
tao model classification_tf2 prune -e /workspace/spec.yaml
        

After the model has been pruned, there might be a slight decrease in accuracy. This happens because some previously useful weights may have been removed. In order to regain the accuracy, NVIDIA recommends that you retrain this pruned model over the same dataset. To do this, use the tao model classification_tf2 train command as documented in Training the model, with an updated spec file that points to the newly pruned model as the pretrained model file.

Users are advised to turn off the regularizer in the train config for classification to recover the accuracy when retraining a pruned model. You may do this by setting the regularizer type to NO_REG. All the other parameters may be retained in the spec file from the previous training.

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.

Here’s an example of the tao model classification_tf2 export command:

            

            
tao model classification_tf2 export [-h] -e <experiment_spec_file> [--gpu_index <gpu_index>]
        

Required Arguments

• -e, --experiment_spec: Path to the spec file.

Optional Arguments

• --gpu_index: The index of (discrete) GPUs used for exporting the model. We can specify the GPU index to run export if the machine has multiple GPUs installed. Note that export can only run on a single GPU.

Sample Usage

Here’s a sample command.

            

            
tao model classification_tf2 export -e /workspace/spec.yaml
        

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.