多任务图像分类

多任务分类需要一个图像目录和两个 CSV 文件，分别用于训练标签和验证标签。图像目录应包含用于训练和验证的所有图像（但可以包含其他图像）。只有训练 CSV 文件中指定的图像将用于训练，验证也是如此。

数据结构应如下所示

            

            
|--dataset_root:
    |--images
        |--1.jpg
        |--2.jpg
        |--3.jpg
        |--4.jpg
        |--5.jpg
        |--6.jpg
    |--train.csv
    |--val.csv
        

训练和验证 CSV 文件包含训练和验证图像的标签。两个 CSV 文件应具有相同的格式：CSV 的第一列必须是 fname，代表图像的文件名。如果您有 N 个任务，则需要额外的 N 列，每列的任务名称作为列名。对于每个图像（CSV 中的行条目），每个任务单元格必须有且只有一个标签。以下是包含 3 个分类任务（颜色、类型和尺寸）的 train.csv 示例

注意：目前，多任务图像分类仅支持 RGB 训练。训练后的模型将始终具有 3 个输入通道。对于灰度图像的推理，用户应将图像作为 RGB 加载，并在所有通道中使用相同的值。这也是训练脚本处理灰度训练图像的方式。

以下是多任务分类的规范文件示例

            

            
random_seed: 42
model_config {
  arch: "resnet"
  n_layers: 101
  use_batch_norm: True
  use_bias: False
  all_projections: False
  use_pooling: True
  use_imagenet_head: True
  resize_interpolation_method: BICUBIC
  input_image_size: "3,224,224"
}
training_config {
  batch_size_per_gpu: 16
  checkpoint_interval: 10
  num_epochs: 80
  enable_qat: false
  learning_rate {
    soft_start_annealing_schedule {
      min_learning_rate: 5e-5
      max_learning_rate: 2e-2
      soft_start: 0.15
      annealing: 0.8
    }
  }
  regularizer {
    type: L1
    weight: 3e-5
  }
  optimizer {
    adam {
      epsilon: 1e-7
      beta1: 0.9
      beta2: 0.999
      amsgrad: false
    }
  }
  pretrain_model_path: "EXPERIMENT_DIR/resnet_101.hdf5"
}
dataset_config {
  image_directory_path: "EXPERIMENT_DIR/data/images"
  train_csv_path: "EXPERIMENT_DIR/data/train.csv"
  val_csv_path: "EXPERIMENT_DIR/data/val.csv"
}
        

模型配置

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

BatchNormalization 参数

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

激活函数

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

训练配置

训练配置 (training_config) 定义训练、评估和推理所需的参数。详细信息总结在下表中。

数据集配置

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

            

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

必需参数

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

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

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

可选参数

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

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

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

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

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

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

            

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

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

multitask_classification 应用程序计算每个任务的评估损失和准确率作为指标。

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

            

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

必需参数

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

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

• -m, --model：提供训练模型的路径。

可选参数

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

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

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

在标记数据集上推理模型可以提供混淆矩阵，您可以从中查看模型在哪里犯错。

TAO 提供了一个命令，可以轻松地为所有任务生成混淆矩阵

            

            
tao model multitask_classification confmat  -i <img_root>
                                      -l <target_csv>
                                      -k <key>
                                      -m <model>
                                      [--gpu_index <gpu_index>]
                                      [--log_file <log_file>]
                                      [-h]
        

必需参数

• -i, --img_root：图像目录的路径。

• -l, --target_csv：真实标签 CSV 文件的路径。

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

• -m, --model：提供训练模型的路径。

可选参数

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

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

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

tao model multitask_classification inference 命令在指定的图像上运行推理。

在 TAO 上训练的多任务分类模型上执行 tao model multitask_classification inference。

            

            
tao model multitask_classification inference -m <model> -i <image> -k <key> -cm <classmap> [--gpu_index <gpu_index>] [--log_file <log_file>] [-h]
        

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

必需参数

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

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

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

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

可选参数

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

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

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

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

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

            

            
tao model multitask_classification prune -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>]
                                   [-h]
        

必需参数

• -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。

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

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

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

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

导出模型

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

            

            
tao model multitask_classification export
            -m <path to the .tlt model file generated by training>
            -k <key>
            -cm <classmap>
            [-o <path to output file>]
            [--gpu_index <gpu_index>]
            [--log_file <log_file_path>]
        

必需参数

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

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

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

可选参数

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

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

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

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

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

有关使用 DeepStream 部署分类模型的更多信息，请参阅 集成多任务图像分类模型 页面。