图像分类 (TF1)

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

以下是模型分类的规范文件示例

            

            
model_config {
  # Model Architecture can be chosen from:
  # ['resnet', 'vgg', 'googlenet', 'alexnet']
  arch: "resnet"
  # for resnet --> n_layers can be [10, 18, 50]
  # for vgg --> n_layers can be [16, 19]
  n_layers: 101
  use_batch_norm: True
  use_bias: False
  all_projections: False
  use_pooling: True
  retain_head: True
  resize_interpolation_method: BICUBIC
  # if you want to use the pretrained model,
  # image size should be "3,224,224"
  # otherwise, it can be "3, X, Y", where X,Y >= 16
  input_image_size: "3,224,224"
}
train_config {
  train_dataset_path: "/path/to/your/train/data"
  val_dataset_path: "/path/to/your/val/data"
  pretrained_model_path: "/path/to/your/pretrained/model"
  # Only ['sgd', 'adam'] are supported for optimizer
  optimizer {
      sgd {
      lr: 0.01
      decay: 0.0
      momentum: 0.9
      nesterov: False
      }
  }
  batch_size_per_gpu: 50
  n_epochs: 150
  # Number of CPU cores for loading data
  n_workers: 16
  # regularizer
  reg_config {
      # regularizer type can be "L1", "L2" or "None".
      type: "L2"
      # if the type is not "None",
      # scope can be either "Conv2D" or "Dense" or both.
      scope: "Conv2D,Dense"
      # 0 < weight decay < 1
      weight_decay: 0.000015
  }
  # learning_rate
  lr_config {
      cosine {
      learning_rate: 0.04
      soft_start: 0.0
      }
  }
  enable_random_crop: True
  enable_center_crop: True
  enable_color_augmentation: True
  mixup_alpha: 0.2
  label_smoothing: 0.1
  preprocess_mode: "caffe"
  image_mean {
    key: 'b'
    value: 103.9
  }
  image_mean {
    key: 'g'
    value: 116.8
  }
  image_mean {
    key: 'r'
    value: 123.7
  }
}
eval_config {
  eval_dataset_path: "/path/to/your/test/data"
  model_path: "/workspace/tao-experiments/classification/weights/resnet_080.tlt"
  top_k: 3
  batch_size: 256
  n_workers: 8
  enable_center_crop: True
}
        

分类实验规范由三个主要组件组成

• model_config

• eval_config

• train_config

模型配置

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

BatchNormalization 参数

参数 batch_norm_config 定义模型中 BatchNormalization 层的参数（动量和 epsilon）。

激活函数

参数 activation 定义模型中激活函数的参数。

Eval Config

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

Training Config

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

学习率调度器

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

Step

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

Soft Annealing

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

Cosine

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

优化器

支持三种类型的优化器：Adam、SGD 和 RMSProp。规范文件中应仅指定一种类型。无论选择哪种类型，它都将包装在 optimizer proto 中，如下例所示

            

            
optimizer {
    sgd {
    lr: 0.01
    decay: 0.0
    momentum: 0.9
    nesterov: False
    }
}
        

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

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

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

使用 tao model classification_tf1 train 命令来调整预训练模型

            

            
tao model classification_tf1 train [-h] -e <spec file>
                              -k <encoding key>
                              -r <result directory>
                              [--gpus <num GPUs>]
                              [--num_processes <number_of_processes>]
                              [--gpu_index <gpu_index>]
                              [--use_amp]
                              [--log_file <log_file_path>]
        

必需参数

• -r, --results_dir：实验输出应写入的文件夹路径。

• -k, --key：用户特定的编码密钥，用于保存或加载 .tlt 模型。

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

可选参数

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

• --num_processes, -np：要为训练生成的进程数。默认值为 -1（等于 --gpus，用于数据并行性用例）。在模型并行性的情况下，应根据实际情况显式将此参数设置为 1 或更大。将 --gpus 设置为大于 1，并将 --num_processes 设置为 1 对应于模型并行性用例；而将 --gpus 和 num_processes 都设置为大于 1 对应于启用模型并行性和数据并行性的情况。例如，--gpus=4 和 --num_processes=2 表示将生成 2 个 Horovod 进程，并且每个进程将占用 2 个 GPU 用于模型并行性。

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

• --use_amp：用于启用 AMP 训练的标志。

• --log_file：日志文件的路径。默认为 stdout。

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

输入要求

• 输入大小：3 * H * W (W, H >= 32)

• 输入格式：JPG、JPEG、PNG

使用示例

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

            

            
tao model classification_tf1 train -e /workspace/tlt_drive/spec/spec.cfg -r /workspace/output -k $YOUR_KEY
        

模型并行性

图像分类支持模型并行性。模型并行性是一种将整个模型拆分到多个 GPU 上的技术，每个 GPU 将保存模型的一部分。模型按层拆分。例如，如果模型有 100 层，那么我们可以将第 0-49 层放在 GPU 0 上，将第 50-99 层放在 GPU 1 上。当模型非常庞大，即使批次大小为 1 也无法容纳在单个 GPU 中时，模型并行性将非常有用。如果我们想增加 BatchNormalization 层看到的批次大小，从而可能提高准确率，模型并行性也很有用。可以通过在 training_config 中设置 model_parallelism 来启用此功能。例如，

            

            
model_parallelism: 0.3
model_parallelism: 0.7
        

将启用 2-GPU 模型并行性，其中第一个 GPU 将保存 30% 的模型层，第二个 GPU 将保存 70% 的模型层。模型层的百分比可以通过一些试错法进行调整，以便所有 GPU 消耗几乎相同的 GPU 内存大小，在这种情况下，我们可以为此模型并行化训练使用最大的批次大小。

模型并行性可以与数据并行性联合启用。例如，在上述情况下，我们启用了 2-GPU 模型并行性，同时我们还可以为其启用 4 个 Horovod 进程。在这种情况下，我们有 4 个用于数据并行性的 Horovod 进程，并且每个进程都将在 2 个 GPU 上拆分模型。

在模型训练完成后，使用实验配置文件，并按照训练模型的步骤操作后，下一步是在测试集上评估此模型，以衡量模型的准确率。TAO 包括 tao model classification_tf1 evaluate 命令来执行此操作。

分类应用程序计算评估损失、Top-k 准确率、精确率和召回率作为指标。

训练完成后，模型将存储在您选择的 $OUTPUT_DIR 输出目录中。使用 tao model classification_tf1 evaluate 命令评估模型

            

            
tao model classification_tf1 evaluate [-h] -e <experiment_spec_file>
                                 -k <key>
                                 [--gpu_index <gpu_index>]
                                 [--log_file <log_file>]
        

必需参数

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

• -k, –key：提供加密密钥以解密模型。

可选参数

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

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

• --log_file：日志文件的路径。默认为 stdout。

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

            

            
tao model classification_tf1 evaluate -e classification_spec.cfg -k $YOUR_KEY
        

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

• 损失

• Top-K 准确率

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

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

• 混淆矩阵

tao model classification_tf1 inference 命令对指定的输入图像集运行推理。对于分类，tao model classification_tf1 inference 为单个图像或 csv 文件（包含图像路径和多个图像的相应标签）提供命令行上的类标签输出。还可以启用 TensorRT Python 推理。

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

            

            
tao model classification_tf1 inference [-h] -m <model>
                                  -i <image>
                                  -d <image  dir>
                                  -k <key>
                                  -cm <classmap>
                                  -e <experiment_spec_file>
                                  [-b <batch size>]
                                  [--gpu_index <gpu_index>]
                                  [--log_file <log_file>]
        

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

必需参数

• -m, --model：预训练模型（TAO 模型）的路径。

• -i, --image：用于推理的单个图像文件。

• -d, --image_dir：用于推理的输入图像目录。

• -k, --key：用于加载模型的密钥。

• -cm, --class_map：指定类索引和标签映射的 json 文件。

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

可选参数

• --batch_size：推理批次大小，默认值：1

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

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

• --log_file：日志文件的路径。默认为 stdout。

剪枝从模型中删除参数以减小模型大小，同时使用 tao model classification_tf1 prune 命令确保模型本身的完整性。

tao model classification_tf1 prune 命令包括以下参数

            

            
tao model classification_tf1 prune [-h] -m <model>
                              -o <output_file>
                              -k <key>
                              [-n <normalizer>
                              [-eq <equalization_criterion>]
                              [-pg <pruning_granularity>]
                              [-pth <pruning threshold>]
                              [-nf <min_num_filters>]
                              [-el <excluded_list>]
                              [--gpu_index <gpu_index>]
                              [--log_file <log_file>]
        

必需参数

• -m, --model：预训练模型的路径

• -o, --output_file：输出检查点的路径

• -k, --key：用于加载 .tlt 模型的密钥

可选参数

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

• -n, –normalizer：max，通过将每层中的每个范数除以最大范数来进行归一化；L2，通过除以包含所有内核范数的向量的 L2 范数来进行归一化。（默认值：max）

• -eq, --equalization_criterion：用于均衡元素级运算层或深度可分离卷积层的输入统计信息的标准。此参数对于 ResNet 和 MobileNet 很有用。选项包括 arithmetic_mean、geometric_mean、union 和 intersection。（默认值：union）

• -pg, -pruning_granularity：一次删除的过滤器数量（默认值：8）

• -pth：用于比较归一化范数的阈值（默认值：0.1）

• -nf, --min_num_filters：每层要保留的最小过滤器数量（默认值：16）

• -el, --excluded_layers：excluded_layers 列表。示例：-i item1 item2（默认值：[]）

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

• --log_file：日志文件的路径。默认为 stdout。

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

使用 Prune 命令

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

            

            
tao model classification_tf1 prune -m /workspace/output/weights/resnet_003.tlt
                         -o /workspace/output/weights/resnet_003_pruned.tlt
                         -eq union
                         -pth 0.7 -k $KEY
        

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

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

导出模型将训练过程与推理分离，并允许在 TAO 环境之外转换为 TensorRT 引擎。TensorRT 引擎特定于每个硬件配置，应针对每个唯一的推理环境生成。导出的模型可以在训练和部署硬件之间通用使用。导出的模型格式称为 .etlt。与 .tlt 一样，.etlt 模型格式也是一种加密模型格式，其密钥与从中导出的 .tlt 模型相同。部署此模型时需要此密钥。

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

            

            
tao model classification_tf1 export [-h] -m <path to the .tlt model file generated by training>
                               -e <experiment_spec_file>
                               [-k <key>]
                               [-o <path to output file>]
                               [--gen_ds_config <Flag to generate ds config and label file>]
                               [--verbose]
                               [--gpu_index <gpu_index>]
                               [--log_file <log_file_path>]
                               [--classmap_json CLASSMAP_JSON]
        

必需参数

• -m, --model：要导出的 .tlt 模型文件的路径。

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

可选参数

• -k, --key：用于保存 .tlt 模型文件的密钥。

• -o, --output_file：保存导出模型的路径。默认为 ./<input_file>.etlt。

• --gen_ds_config：一个布尔标志，指示是否在与 output_file 相同的目录中生成模板 DeepStream 相关配置 (“nvinfer_config.txt”) 以及标签文件 (“labels.txt”)。请注意，配置文件不是完整的配置文件，需要用户使用生成的参数更新 DeepStream 中的示例配置文件。

• --classmap_json：classmap_json 文件的路径。它已在训练结果文件夹中生成。如果启用了 gen_ds_config，则需要此文件。

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

• --log_file：日志文件的路径。默认为 stdout。

• -v, --verbose：详细日志。

使用示例

以下是一个示例命令。

            

            
tao model classification_tf1 export -m /ws/output_retrain/weights/resnet_001.tlt
                              -o /ws/export/final_model.etlt
                              -k $KEY
        

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

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