图像分类 (TF2)

有关图像分类的数据格式的更多信息，请参阅数据标注格式页面。

以下是分类规范文件的示例。它有七个主要组成部分——dataset、model、train、evaluate、inference、prune 和 export——以及加密密钥 (encryption_key) 和结果目录 (results_dir) 的参数。

            

            
results_dir: '/workspace/results_dir'
encryption_key: 'nvidia_tlt'
num_gpus: 1
gpu_ids: [0]
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'
        

规范文件的格式为 YAML。下表总结了规范文件的顶层结构

模型配置

下表描述了 model 配置中的参数。

数据集配置

下表描述了 data 配置中的参数。

下表描述了 dataset.augmentation 配置中的参数。

推理配置

下表定义了用于评估分类模型的可配置参数。

评估配置

下表定义了用于评估分类模型的可配置参数。

训练配置

本节定义了分类模型训练器的可配置参数。

学习率调度器

参数 lr_config 定义了学习率调度器的参数。学习率调度器可以是 step、soft_anneal 或 cosine。

Step

参数 step 定义了步进学习率调度器。

软退火

参数 soft_anneal 定义了软退火学习率调度器。

Cosine

参数 cosine 定义了余弦学习率调度器。

优化器配置

支持三种类型的优化器：Adam、SGD 和 RMSProp。规范文件中应仅使用 optimizer 指定一种类型

下表总结了 Adam 优化器参数。

下表总结了 SGD 优化器参数。

下表总结了 RMSProp 优化器参数。

BatchNorm 配置

剪枝配置

prune 配置定义了训练模型的剪枝过程。下表总结了详细描述。

导出配置

export 配置包含将 .tlt 模型导出到 .onnx 模型的参数，该模型可用于部署。

使用 tao model classification_tf2 train 命令微调预训练模型

            

            
tao model classification_tf2 train [-h] -e <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>]
                             [num_gpus=<num GPUs>]
                             [gpu_ids=<gpu_index>]
        

必需参数

• -e, --experiment_spec_file：实验规范文件的路径。

可选参数

• model.<model_option>：模型选项。

• dataset.<dataset_option>：数据集选项。

• train.<train_option>：训练选项。

• num_gpus：用于训练和启动进程的 GPU 数量。默认值为 1。

• gpu_ids：用于运行训练的 GPU 索引。当机器安装了多个 GPU 时，我们可以指定用于运行训练的 GPU 索引。

• -h, --help：打印帮助消息。

输入要求

• 输入大小：3 * 高 * 宽（宽、高 >= 32）

• 输入格式：JPG、JPEG、PNG

示例用法

以下是使用 tao model classification_tf2 train 命令的示例

            

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

在模型经过训练后，使用实验配置文件，并按照训练模型的步骤操作，下一步是在测试集上评估此模型，以衡量模型的准确性。

分类应用程序计算评估损失、Top-k 准确率、精确率和召回率作为指标。使用 tao model classification_tf2 evaluate 命令评估模型

            

            
tao model classification_tf2 evaluate [-h] -e <experiment_spec_file>
                             evaluate.checkpoint=<model to be evaluated>
                             [evaluate.<evaluate_option>=<evaluate_option_value>]
                             [gpu_ids=<gpu_index>]
        

必需参数

• -e, --experiment_spec_file：实验规范文件的路径。

• evaluate.checkpoint：要评估的 .pth 模型。

可选参数

• -h, --help：显示此帮助消息并退出。

• gpu_ids：用于运行评估的 GPU 索引。仅支持 1 个 GPU。

• evaluate.<evaluate_option>：评估选项。

如果您按照训练分类模型中的示例操作，请运行评估

            

            
tao model classification_tf2 evaluate -e classification_spec.yaml
        

TAO 评估分类并生成以下指标

• 损失

• Top-K 准确率

• 精确率 (P)：TP / (TP + FP)

• 召回率 (R)：TP / (TP + FN)

• 混淆矩阵

tao model classification_tf2 inference 命令在一组指定的输入图像上运行推理。对于分类，tao model classification_tf2 inference 通过命令行为单个图像或包含图像路径以及多个图像的相应标签的 csv 文件提供类别标签输出。还可以启用 TensorRT Python 推理。

在 TAO 上训练的分类模型上执行 tao model classification_tf2 inference。

            

            
tao model classification_tf2 inference [-h] -e <experiment_spec_file>
                             inference.checkpoint=<model to be inferenced>
                             [inference.<inference_option>=<inference_option_value>]
                             [gpu_ids=<gpu_index>]
        

以下是 tao model classification_tf2 inference 工具的参数

必需参数

• -e, --experiment_spec_file：实验规范文件的路径。

• inference.checkpoint：要推理的 .pth 模型。

可选参数

• -h, --help：显示此帮助消息并退出。

• gpu_ids：用于运行推理的 GPU 索引。仅支持 1 个 GPU。

• inference.<inference_option>：推理选项。

剪枝从模型中移除参数，以减小模型大小，同时使用 tao model classification_tf2 prune 命令，而不损害模型本身的完整性。

tao model classification_tf2 prune 命令包含以下参数

            

            
tao model classification_tf2 prune [-h] -e <experiment_spec_file>
                             prune.checkpoint=<model to be pruned>
                             [prune.<prune_option>=<prune_option_value>]
        

必需参数

• -e, --experiment_spec_file：实验规范文件的路径。

• prune.checkpoint：要剪枝的 .pth 模型。

可选参数

• -h, --help：显示此帮助消息并退出。

• gpu_ids：用于运行剪枝的 GPU 索引。仅支持 1 个 GPU。

• prune.<prune_option>：剪枝选项。

剪枝后，需要重新训练模型。有关更多详细信息，请参阅重新训练剪枝模型。

使用剪枝命令

以下是使用 tao model classification_tf2 prune 命令的示例

            

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

模型剪枝后，准确率可能会略有下降。发生这种情况是因为某些以前有用的权重可能已被移除。为了重新获得准确率，NVIDIA 建议您在同一数据集上重新训练此剪枝模型。为此，请使用 训练模型中记录的 tao model classification_tf2 train 命令，并使用更新后的规范文件，该文件指向新剪枝的模型作为预训练模型文件。

建议用户在分类的 train 配置中关闭正则化器，以在重新训练剪枝模型时恢复准确率。您可以通过将正则化器 type 设置为 NO_REG 来执行此操作。所有其他参数都可以从之前的训练中保留在规范文件中。

导出模型将训练过程与推理分离，并允许在 TAO 环境之外转换为 TensorRT 引擎。TensorRT 引擎特定于每个硬件配置，应为每个唯一的推理环境生成。导出的模型可以在训练和部署硬件上通用使用。

以下是 tao model classification_tf2 export 命令的示例

            

            
tao model classification_tf2 export [-h] -e <experiment_spec_file>
                             export.checkpoint=<model to export>
                             export.onnx_file=<onnx path>
                             [export.<export_option>=<export_option_value>]
                             [gpu_ids=<gpu_index>]
        

必需参数

• -e, --experiment_spec：规范文件的路径。

• export.checkpoint：要导出的 .pth 模型。

• export.onnx_file：保存 .etlt 或 .onnx 模型的路径。

可选参数

• gpu_ids：用于导出模型的（离散）GPU 的索引。如果机器安装了多个 GPU，我们可以指定 GPU 索引来运行导出。请注意，导出只能在单个 GPU 上运行。

• export.<export_option>：导出选项。

示例用法

以下是示例命令。

            

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

对于 TensorRT 引擎生成、验证和 int8 校准，请参阅 TAO 部署文档。

有关使用 DeepStream 部署分类模型的更多信息，请参阅集成分类 (TF1/TF2/PyTorch) 模型页面。