RetinaNet

RetinaNet 支持以下任务

• dataset_convert

• train

• evaluate

• prune

• inference

• export

这些任务可以通过命令行使用 TAO 启动器调用，遵循以下提及的约定

            

            
tao model retinanet <sub_task> <args_per_subtask>
        

其中，args_per_subtask 是给定子任务所需的命令行参数。以下详细解释了每个子任务。

TAO 中的对象检测应用程序期望使用 KITTI 格式的数据进行训练和评估。

有关 KITTI 数据格式的更多信息，请参阅数据注释格式页面。

RetinaNet 数据加载器支持原始 KITTI 格式的数据以及 TFrecords。

要使用 TFRecords 优化跨数据批次的迭代，首先需要将原始输入数据转换为 TFRecords 格式。这可以使用 dataset_convert 子任务完成。目前，支持 KITTI 和 COCO 格式。

dataset_convert 工具需要一个配置文件作为输入。以下章节包含配置文件的详细信息和示例。

数据集转换器的配置文件

dataset_convert 工具提供了几个可配置的参数。这些参数封装在一个 spec 文件中，用于将数据从原始注释格式转换为训练器可以摄取的 TFRecords 格式。KITTI 和 COCO 格式可以通过分别使用 kitti_config 或 coco_config 进行配置。在一个 spec 文件中，您只能使用其中之一。Spec 文件是一个 prototxt 格式的文件，具有以下全局参数

• kitti_config：一个嵌套的 prototxt 配置，具有多个输入参数

• coco_config：一个嵌套的 prototxt 配置，具有多个输入参数

• image_directory_path：数据集根目录的路径。image_dir_name 将附加到此路径以获取输入图像，并且必须与实验 spec 文件中指定的路径相同。

• target_class_mapping：prototxt 字典，用于将 tfrecords 中的类名称映射到网络中要训练的目标类。

kitti_config

以下是 kitti_config 字段的可配置参数的描述

下面显示的示例配置文件将 100% KITTI 数据集转换为训练集。

            

            
kitti_config {
  root_directory_path: "/workspace/tao-experiments/data/"
  image_dir_name: "training/image_2"
  label_dir_name: "training/label_2"
  image_extension: ".png"
  partition_mode: "random"
  num_partitions: 2
  val_split: 0
  num_shards: 10
}
image_directory_path: "/workspace/tao-experiments/data/"
  target_class_mapping {
      key: "car"
      value: "car"
  }
  target_class_mapping {
      key: "pedestrian"
      value: "pedestrian"
  }
  target_class_mapping {
      key: "cyclist"
      value: "cyclist"
  }
  target_class_mapping {
      key: "van"
      value: "car"
  }
  target_class_mapping {
      key: "person_sitting"
      value: "pedestrian"
  }
  target_class_mapping {
      key: "truck"
      value: "car"
  }
        

coco_config

以下是 coco_config 字段的可配置参数的描述

下面显示的示例配置文件转换包含训练和验证数据的 COCO 数据集，其中验证的分片数为 32，训练的分片数为 256。

            

            
coco_config {
  root_directory_path: "/workspace/tao-experiments/data/coco"
  image_dir_names: ["val2017", "train2017"]
  annotation_files: ["annotations/instances_val2017.json", "annotations/instances_train2017.json"]
  num_partitions: 2
  num_shards: [32, 256]
}
image_directory_path: "/workspace/tao-experiments/data/coco"
        

数据集转换器工具的示例用法

下面描述了 dataset_convert 工具

            

            
tao model retinanet dataset-convert [-h] -d DATASET_EXPORT_SPEC
                                   -o OUTPUT_FILENAME
                                   [-v]
        

您可以使用以下参数

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

• -d, --dataset-export-spec：检测数据集 spec 的路径，其中包含导出 .tfrecord 文件的配置

• -o, --output_filename：输出文件名

• -v：启用详细模式以显示调试消息

以下示例显示了如何将命令与数据集一起使用

            

            
tao model retinanet dataset_convert -d /path/to/spec.txt
                              -o /path/to/tfrecords/train
        

以下是 RetinaNet spec 文件的示例。它有 6 个主要组成部分：retinanet_config、training_config、eval_config、nms_config、augmentation_config 和 dataset_config。Spec 文件的格式是 protobuf 文本 (prototxt) 消息，其每个字段可以是基本数据类型或嵌套消息。下表总结了 spec 文件的顶层结构

            

            
random_seed: 42
retinanet_config {
  aspect_ratios_global: "[1.0, 2.0, 0.5]"
  scales: "[0.045, 0.09, 0.2, 0.4, 0.55, 0.7]"
  two_boxes_for_ar1: false
  clip_boxes: false
  loss_loc_weight: 0.8
  focal_loss_alpha: 0.25
  focal_loss_gamma: 2.0
  variances: "[0.1, 0.1, 0.2, 0.2]"
  arch: "resnet"
  nlayers: 18
  n_kernels: 1
  n_anchor_levels: 1
  feature_size: 256
  freeze_bn: false
  freeze_blocks: 0
}
training_config {
  enable_qat: False
  batch_size_per_gpu: 24
  num_epochs: 100
  pretrain_model_path: "YOUR_PRETRAINED_MODEL"
  optimizer {
    sgd {
      momentum: 0.9
      nesterov: True
    }
  }
  learning_rate {
  soft_start_annealing_schedule {
    min_learning_rate: 4e-5
    max_learning_rate: 1.5e-2
    soft_start: 0.15
    annealing: 0.5
    }
  }
  regularizer {
    type: L1
    weight: 2e-5
  }
}
eval_config {
  validation_period_during_training: 10
  average_precision_mode: SAMPLE
  batch_size: 24
  matching_iou_threshold: 0.5
}
nms_config {
  confidence_threshold: 0.01
  clustering_iou_threshold: 0.6
  top_k: 200
}
augmentation_config {
    output_width: 384
    output_height: 1248
    output_channel: 3
    image_mean {
      key: 'b'
      value: 103.9
    }
    image_mean {
      key: 'g'
      value: 116.8
    }
    image_mean {
      key: 'r'
      value: 123.7
    }
}
dataset_config {
  data_sources: {
      # option 1
      tfrecords_path: "/workspace/tao-experiments/data/tfrecords/kitti_train*"

      # option 2
      # label_directory_path: "/workspace/tao-experiments/data/training/label_2"
      # image_directory_path: "/workspace/tao-experiments/data/training/image_2"
  }
  target_class_mapping {
      key: "car"
      value: "car"
  }
  target_class_mapping {
      key: "pedestrian"
      value: "pedestrian"
  }
  target_class_mapping {
      key: "cyclist"
      value: "cyclist"
  }
  target_class_mapping {
      key: "van"
      value: "car"
  }
  target_class_mapping {
      key: "person_sitting"
      value: "pedestrian"
  }
  validation_data_sources: {
      label_directory_path: "/workspace/tao-experiments/data/val/label"
      image_directory_path: "/workspace/tao-experiments/data/val/image"
  }
}
        

训练配置

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

训练可视化配置

训练期间的可视化通过 visualizer 参数配置，该参数包含以下描述的参数。

如果启用了可视化，则将在训练期间生成 tensorboard 日志，包括学习率、训练损失、验证损失、验证 mAP 和每个类的验证 AP 的图表。带有 bbox 的增强图像也将在 tensorboard 中生成。

提前停止

下表描述了提前停止的参数。

评估配置

评估配置 (eval_config) 定义了在训练期间或独立评估所需的参数。下表总结了详细信息。

NMS 配置

NMS 配置 (nms_config) 定义了 NMS 后处理所需的参数。NMS 配置适用于训练、验证、评估、推理和导出中模型的 NMS 层。下表总结了详细信息。

增强配置

augmentation_config 参数定义了预处理后的图像大小。SSD 论文中的增强方法将在训练期间执行，包括随机翻转、放大、缩小和颜色抖动。增强后的图像将调整为 augmentation_config 中定义的输出形状。在评估过程中，仅执行调整大小。

数据集配置

RetinaNet 数据加载器假定数据以 KITTI 格式准备（图像和标签在两个单独的文件夹中，其中图像文件夹中的每个图像在标签文件夹中都有一个具有相同文件名的 txt 标签文件。标签文件内容遵循 KITTI 格式），并且已完成训练/验证拆分。

dataset_config 中的参数定义如下

• data_sources：捕获要训练的数据集的路径。如果您有多个训练数据源，则可以使用多个 data_sources。此字段包含 3 个参数：* label_directory_path：数据源标签文件夹的路径 * image_directory_path：数据源图像文件夹的路径 * tfrecords_path：TFRecords 的路径

  当使用原始 KTTTI 格式的数据作为输入时，仅需要 label_directory_path 和 image_directory_path。当使用 TFRecords 作为数据摄取时，仅需要 tfrecords_path。

• include_difficult_in_training：是否在训练中包含困难框。如果设置为 false，则将忽略困难框。困难框是 KITTI 标签中遮挡级别为 2 的框。（仅适用于原始 KITTI 格式的数据）

• target_class_mapping：此参数将标签中的类名称映射到网络中要训练的目标类。为每个源类到目标类的映射定义一个元素。包含此字段的目的是将相似的类对象归类到一个保护伞下。例如：汽车、面包车、重型卡车等可以归类为汽车。“key” 字段是 tfrecords 文件中类名称的值，“value” 字段对应于网络预期学习的值。

• validation_data_sources：捕获要验证的数据集的路径。如果您有多个验证数据源，则可以使用多个 validation_data_sources。此字段包含 2 个参数

  • label_directory_path：数据源标签文件夹的路径

  • image_directory_path：数据源图像文件夹的路径

类权重配置

RetinaNet 支持在训练期间对损失函数进行类级别加权。以下是类权重配置 (class_weighting_config) 的示例，用于设置两个类的权重。

            

            
class_weighting_config {
  class_weighting{
      key: "person"
      value: 1.0
  }
  class_weighting{
      key: "bus"
      value: 5.0
  }
}
        

class_weighting_config 中的参数定义如下

• class_weighting：此参数将类名称映射到损失函数上相应的类权重。类权重 value 应大于 0。如果未在配置中显式设置类权重，则将隐式分配默认值 1.0。

RetinaNet 配置

RetinaNet 配置 (retinanet_config) 定义了构建 RetinaNet 模型所需的参数。下表总结了详细信息。

Focal loss 的计算方式如下

方差

使用以下命令训练 RetinaNet 模型

            

            
tao model retinanet train [-h] -e <experiment_spec>
                         -r <output_dir>
                         -k <key>
                         [--gpus <num_gpus>]
                         [--gpu_index <gpu_index>]
                         [--use_amp]
                         [--log_file <log_file_path>]
        

必需参数

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

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

• -e, --experiment_spec_file：用于设置评估实验的实验规范文件。这应与训练规范文件相同。

可选参数

• --gpus：在多 GPU 场景中用于训练的 GPU 数量（默认值：1）。

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

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

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

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

输入要求

• 输入大小：C * W * H（其中 C = 1 或 3，W >= 128，H >= 128，W、H 是 32 的倍数）

• 图像格式：JPG、JPEG、PNG

• 标签格式：KITTI 检测

示例用法

以下是在 RetinaNet 模型上使用 train 命令的示例

            

            
tao model retinanet train --gpus 2 -e /path/to/spec.txt -r /path/to/result -k $KEY
        

要对 RetinaNet 模型运行评估，请使用以下命令

            

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

必需参数

• -e, --experiment_spec_file：用于设置评估实验的实验 spec 文件。这应与训练规范文件相同。

• -m, --model：用于评估的模型文件的路径（同时支持 TAO 模型和 TensorRT 引擎）。

• -k, --key：提供密钥以加载 TAO 模型（如果使用 TensorRT 引擎，则不需要密钥）。

可选参数

• --gpu_index：用于运行评估的 GPU 索引。当机器安装了多个 GPU 时，我们可以指定用于运行评估的 GPU 索引。请注意，评估只能在单个 GPU 上运行。

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

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

示例用法

以下是在 RetinaNet 模型上使用 evaluate 命令的示例

            

            
tao model retinanet evaluate -e /path/to/spec.txt -m /path/to/model.tlt -k $KEY
        

RetinaNet 网络的推理工具可用于可视化 bbox，或在图像目录上生成逐帧 KITTI 格式标签。支持两种模式，即 TAO 模型模式和 TensorRT 引擎模式。您可以使用以下命令执行 TAO 模型模式

            

            
tao model retinanet inference [-h] -i <input directory>
                             -o <output annotated image directory>
                             -e <experiment spec file>
                             -m <model file>
                             -k <key>
                             [-l <output label directory>]
                             [-t <visualization threshold>]
                             [--gpu_index <gpu_index>]
                             [--log_file <log_file_path>]
        

必需参数

• -m, --model：预训练模型的路径（同时支持 TAO 模型和 TensorRT 引擎）。

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

• -o, --out_image_dir：输出注释图像的目录路径。

• -k, --key：用于加载 TAO 模型的密钥（如果使用 TensorRT 引擎，则不需要密钥）。

• -e, --config_path：训练的实验 spec 文件的路径。

可选参数

• -t, --threshold：用于绘制 bbox 的阈值。默认值：0.3

• -l, --out_label_dir：用于输出 KITTI 标签的目录。

• --gpu_index：用于在其上运行推理的 GPU 索引。如果机器安装了多个 GPU，我们可以指定用于运行推理的 GPU 索引。请注意，推理只能在单个 GPU 上运行。

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

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

示例用法

这是一个在 RetinaNet 模型上使用推理命令的示例

            

            
tao model retinanet inference -e /path/to/spec.txt -m /path/to/model.tlt -k $KEY
                        -o /path/to/output_dir -i /path/to/input_dir
        

剪枝操作通过移除模型中的参数来减小模型大小，同时使用 tao model retinanet prune 命令确保模型本身的完整性。

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

            

            
tao model retinanet prune [-h] -m <retinanet 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_path>]
        

必需参数

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

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

• -k, --key：加载 :code`.tlt` 模型的密钥。

可选参数

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

  注意

  NVIDIA 建议更改阈值，以使模型中的参数数量保持在原始未剪枝模型的 10-20% 以内。

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

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

• --gpu_index：用于运行剪枝的 GPU 索引。如果机器安装了多个 GPU，我们可以指定用于运行剪枝的 GPU 索引。请注意，剪枝只能在单个 GPU 上运行。

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

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

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

使用剪枝命令

这是一个使用 tao model retinanet prune 命令的示例

            

            
tao model retinanet 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 retinanet train 命令，如 训练模型 中所述，并更新 spec 文件，将新剪枝的模型指定为预训练模型文件。

建议用户在 RetinaNet 的 training_config 中关闭正则化器，以在重新训练剪枝模型时恢复精度。您可以通过将正则化器类型设置为 NO_REG 来执行此操作，如 Training config 中所述。所有其他参数可以从之前的训练中保留在 spec 文件中。

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

INT8 模式概述

TensorRT 引擎可以在 INT8 模式下生成以提高性能，但需要在引擎创建时使用校准缓存。如果使用 --data_type 标志设置为 int8 运行 tao model retinanet export，则使用校准张量文件生成校准缓存。预先生成校准信息并缓存它消除了在推理机器上校准模型的需要。移动校准缓存通常比移动校准张量文件更方便，因为它是一个小得多的文件，可以与导出的模型一起移动。使用校准缓存还可以加快引擎创建速度，因为构建缓存可能需要几分钟才能生成，具体取决于 Tensorfile 的大小和模型本身。

导出工具可以通过使用以下任一选项摄取训练数据来生成 INT8 校准缓存

• 选项 1：使用训练数据加载器加载训练图像以进行 INT8 校准。此选项现在是推荐的方法，通过利用训练数据集加载器来支持多个图像目录。这也确保了校准期间数据的两个重要方面

  • INT8 校准步骤中的数据预处理与训练过程中的相同。

  • 数据批次在整个训练数据集中随机采样，从而提高了 INT8 模型的准确性。

• 选项 2：将工具指向您要用于校准模型的图像目录。对于此选项，请确保创建一个随机图像的子采样目录，以最好地代表您的训练数据集。

FP16/FP32 模型

如果您需要在 INT8 精度下运行推理，则仅需要 calibration.bin。对于基于 FP16/FP32 的推理，导出步骤要简单得多。所需的只是提供来自训练/重新训练步骤的 .tlt 模型，以转换为 .etlt。

导出 RetinaNet 模型

这是一个 tao model retinanet export 命令的命令行参数示例

            

            
tao model retinanet export [-h] -m <path to the .tlt model file>
                          --experiment_spec <path to experiment spec file>
                          -k <key>
                          [-o <path to output file>]
                          [--cal_data_file <path to tensor file>]
                          [--cal_image_dir <path to the directory images to calibrate the model]
                          [--cal_cache_file <path to output calibration file>]
                          [--data_type <Data type for the TensorRT backend during export>]
                          [--batches <Number of batches to calibrate over>]
                          [--max_batch_size <maximum trt batch size>]
                          [--max_workspace_size <maximum workspace size]
                          [--batch_size <batch size to TensorRT engine>]
                          [--engine_file <path to the TensorRT engine file>]
                          [--gen_ds_config]
                          [--strict_type_constraints]
                          [--force_ptq]
                          [--gpu_index <gpu_index>]
                          [--log_file <log_file_path>]
                          [--verbose]
        

必需参数

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

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

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

可选参数

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

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

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

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

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

QAT 导出模式必需的参数

• --cal_json_file：包含 QAT 模型张量比例的 json 文件的路径。如果要生成 QAT 模型的引擎，则此参数是必需的。

使用示例

这是一个导出 RetinaNet 模型的示例命令。

            

            
tao model retinanet export -m /ws/retinanet_resnet18_epoch_100.tlt  \
                         -o /ws/retinanet_resnet18_epoch_100_int8.etlt \
                         -e /ws/retinanet_retrain_resnet18_kitti.txt \
                         -k $KEY
        

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

有关将 RetinaNet 模型部署到 DeepStream 的更多信息，请参阅 集成 RetinaNet 模型 页面。