二、TAO——训练YOLOv3
YOLOv3 是包含在 TAO 工具包中的目标检测模型。YOLOv3 支持以下任务：
- dataset_convert
- kmeans
- train
- evaluate
- inference
- prune
- export
这些任务可以通过在命令行中使用以下格式从 TAO 工具包启动器中调用：
tao model yolo_v3 <sub_task> <args_per_subtask>
其中 args_per_subtask 是针对给定子任务所需的命令行参数。每个子任务将在下面做详细解释。
1 Preparing the Input Data Structure
YOLOv3的数据集结构与DetectNet_v2的相同，唯一的区别在于生成TFRecords的命令行不同。要为YOLOv3训练生成TFRecords，请使用以下命令：
tao model yolo_v3 dataset_convert [-h] -d <dataset_spec>
-o <output_tfrecords_file>
[--gpu_index <gpu_index>]
1.1 Required Arguments
- -d, –dataset_spec：数据集规范文件的路径。
- -o, –output_filename：输出的TFRecords文件的路径。
1.2 Optional Arguments
- –gpu_index：要在哪个GPU上运行此命令的GPU索引。如果机器上安装了多个GPU，可以指定用于运行此命令的GPU索引。请注意，此命令只能在单个GPU上运行。
2 Creating a Configuration File
以下是YOLOv3规范文件的示例。它包括六个主要组件：yolov3_config、training_config、eval_config、nms_config、augmentation_config和dataset_config。规范文件的格式是一个protobuf文本（prototxt）消息，它的每个字段可以是基本数据类型或嵌套消息。规范文件的顶级结构总结在下表中。
random_seed: 42
yolov3_config {
big_anchor_shape: "[(114.94, 60.67), (159.06, 114.59), (297.59, 176.38)]"
mid_anchor_shape: "[(42.99, 31.91), (79.57, 31.75), (56.80, 56.93)]"
small_anchor_shape: "[(15.60, 13.88), (30.25, 20.25), (20.67, 49.63)]"
matching_neutral_box_iou: 0.7
arch: "resnet"
nlayers: 18
arch_conv_blocks: 2
loss_loc_weight: 0.8
loss_neg_obj_weights: 100.0
loss_class_weights: 1.0
freeze_bn: false
force_relu: false
}
training_config {
batch_size_per_gpu: 8
num_epochs: 80
enable_qat: false
checkpoint_interval: 10
learning_rate {
soft_start_annealing_schedule {
min_learning_rate: 1e-6
max_learning_rate: 1e-4
soft_start: 0.1
annealing: 0.5
}
}
regularizer {
type: L1
weight: 3e-5
}
optimizer {
adam {
epsilon: 1e-7
beta1: 0.9
beta2: 0.999
amsgrad: false
}
}
pretrain_model_path: "EXPERIMENT_DIR/pretrained_resnet18/tlt_pretrained_object_detection_vresnet18/resnet_18.hdf5"
}
eval_config {
average_precision_mode: SAMPLE
batch_size: 8
matching_iou_threshold: 0.5
}
nms_config {
confidence_threshold: 0.001
clustering_iou_threshold: 0.5
top_k: 200
}
augmentation_config {
hue: 0.1
saturation: 1.5
exposure:1.5
vertical_flip:0
horizontal_flip: 0.5
jitter: 0.3
output_width: 1248
output_height: 384
output_channel: 3
randomize_input_shape_period: 0
}
dataset_config {
data_sources: {
tfrecords_path: "/workspace/tao-experiments/data/tfrecords/kitti_trainval/kitti_trainval*"
image_directory_path: "/workspace/tao-experiments/data/training"
}
include_difficult_in_training: true
image_extension: "png"
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_fold: 0
}
2.1 Training Config
训练配置（training_config）定义了用于训练、评估和推断所需的参数。下表总结了详细信息。
| 字段名 | 描述 | 数据类型和约束 | 推荐/典型值 |
|---|---|---|---|
| batch_size_per_gpu | 每个GPU的批处理大小，因此有效批处理大小是batch_size_per_gpu * num_gpus | 无符号整数，正数 | - |
| checkpoint_interval | 每个模型检查点/验证运行的训练时代数 | 无符号整数，正数 | 10 |
| num_epochs | 训练网络的时代数 | 无符号整数，正数。 | - |
| enable_qat | 是否使用量化感知训练 | 布尔值 | - |
| learning_rate | 仅支持带有以下嵌套参数的“soft_start_annealing_schedule”： | 消息类型 | - |
| min_learning_rate: 整个实验期间的最低学习率 | |||
| max_learning_rate: 整个实验期间的最大学习率 | |||
| soft_start: 热身前的经过的时间（以0到1之间的进度百分比表示） | |||
| regularizer | 此参数配置在训练期间要使用的正则化器，并包含以下嵌套参数： | 消息类型 | - |
| type: 要使用的正则化器类型。NVIDIA支持NO_REG、L1或L2 | |||
| weight: 正则化器权重，浮点值 | |||
| optimizer | 可以是“adam”、“sgd”或“rmsprop”中的一个。每种类型都具有以下参数： | 消息类型 | - |
| adam: epsilon, beta1, beta2, amsgrad | |||
| sgd: momentum, nesterov | |||
| rmsprop: rho, momentum, epsilon, centered | |||
| pretrain_model_path | 如果有的话，预训练模型的路径。pretrain_model_path、resume_model_path或pruned_model_path中最多只能存在一个。 | 字符串 | - |
| resume_model_path | 如果有的话，要继续训练的TAO检查点模型的路径。pretrain_model_path、resume_model_path或pruned_model_path中最多只能存在一个。 | 字符串 | - |
| pruned_model_path | 如果有的话，要重新训练的TAO剪枝模型的路径。pretrain_model_path、resume_model_path或pruned_model_path中最多只能存在一个。 | 字符串 | - |
| max_queue_size | 数据加载中的预取批次数 | 无符号整数，正数 | - |
| n_workers | 每个GPU的数据加载工作进程数 | 无符号整数，正数 | - |
| use_multiprocessing | 是否使用keras序列数据加载器的多进程模式 | 布尔值 | true (在发生死锁时，重新启动训练并使用False) |
2.2 Evaluation Config
评估配置（eval_config）定义了在训练期间或作为独立过程进行评估所需的参数。以下是详细信息总结在下表中：
| 字段名 | 描述 | 数据类型和约束 | 推荐/典型值 |
|---|---|---|---|
| average_precision_mode | 平均精度（AP）计算模式，可以是SAMPLE或INTEGRATE。 SAMPLE用作VOC 2009或之前的VOC指标。 INTEGRATE用于VOC 2010或之后。 | 枚举类型（SAMPLE或INTEGRATE） | SAMPLE |
| matching_iou_threshold | 被认为是匹配的预测框和地面真实框之间的最低IoU。 | 浮点数 | 0.5 |
2.3 NMS Config
NMS配置（nms_config）定义了用于NMS后处理所需的参数。NMS配置适用于模型的NMS层，在训练、验证、评估、推理和导出过程中都会使用。以下是详细信息总结在下表中：
| 字段名 | 描述 | 数据类型和约束 | 推荐/典型值 |
|---|---|---|---|
| confidence_threshold | 在应用NMS之前，具有置信度得分小于confidence_threshold的框将被丢弃。 | 浮点数 | 0.01 |
| cluster_iou_threshold | 框将通过NMS过程的IoU阈值以下。 | 浮点数 | 0.6 |
| top_k | 在NMS Keras层之后，将输出top_k个框。如果有效框的数量小于k，返回的数组将填充置信度得分为0的框。 | 无符号整数 | 200 |
| infer_nms_score_bits | 在TensorRT OSS中NMS插件中表示得分值的位数。有效范围是[1, 10]的整数。将其设置为其他值将使其回退到普通的NMS。目前，这个优化的NMS插件只在FP16中可用，但应选择INT8数据类型，因为TensorRT OSS中没有INT8 NMS，因此将选择在FP16中的最快实现。如果回退到普通的NMS，则在构建引擎时的实际数据类型将决定要运行的精度（FP16或FP32）。 | 整数，取值范围为[1, 10] | 0 |
2.4 Augmentation Config
增强配置（augmentation_config）定义了在线数据增强所需的参数。详细信息总结在下表中。
| 字段名 | 描述 | 数据类型和约束 | 推荐/典型值 |
|---|---|---|---|
| hue | 图像色调在[-hue, hue] * 180.0范围内变化 | 浮点数，取值范围为[0, 1] | 0.1 |
| saturation | 图像饱和度在[1.0 / saturation, saturation]范围内变化 | 浮点数，大于等于1.0 | 1.5 |
| exposure | 图像曝光度在[1.0 / exposure, exposure]范围内变化 | 浮点数，大于等于1.0 | 1.5 |
| vertical_flip | 图像垂直翻转的概率 | 浮点数，取值范围为[0, 1] | 0 |
| horizontal_flip | 图像水平翻转的概率 | 浮点数，取值范围为[0, 1] | 0.5 |
| jitter | 增强中允许的最大抖动。此处的抖动是指YOLO网络中的抖动增强 | 浮点数，取值范围为[0, 1] | 0.3 |
| output_width | 增强管道的基本输出图像宽度 | 整数，必须是32的倍数 | - |
| output_height | 增强管道的基本输出图像高度 | 整数，必须是32的倍数 | - |
| output_channel | 增强管道的输出通道数 | 1或3 | - |
| randomize_input_shape_period | 随机更改输出宽度和高度的批处理间隔。对于值K，增强管道将每K批次调整一次输出形状，调整后的输出宽度/高度将在基本宽度/高度的0.6到1.5倍之间。注意：如果K=0，输出宽度/高度将始终与配置的基本宽度/高度匹配，训练速度将更快，但训练网络的准确性可能不如预期。 | 非负整数 | 10 |
| image_mean | 用于指定图像均值的键/值对。如果省略，图像预处理将使用ImageNet均值。如果设置了值，根据output_channel，必须配置“r/g/b”或“l”键/值对之一。 | 字典 | - |
2.5 Dataset Config
YOLOv3支持两种数据格式：序列格式（包含图像文件夹和具有KITTI格式的原始标签文件夹）和tfrecords格式（包含图像文件夹和TFRecords）。在大多数情况下，使用TFRecord数据集进行训练比序列格式更快，因此推荐使用TFRecord数据集的格式。但在某些情况下，如小尺寸输入分辨率（例如416x416），序列格式稍微快于TFRecord。
YOLOv3的数据加载器假设训练/验证分割已经完成，并且数据以KITTI格式准备好：图像和标签分别存储在两个文件夹中，图像文件夹中的每个图像都有一个与标签文件夹中具有相同文件名的.txt标签文件，标签文件的内容遵循KITTI格式。COCO数据格式也受支持，但仅通过TFRecords格式。您可以使用dataset_convert来准备TFRecords数据。
以下是从KITTI数据格式转换为TFRecord数据格式的示例dataset_config。
dataset_config {
data_sources: {
tfrecords_path: "/workspace/tao-experiments/data/tfrecords/kitti_trainval/kitti_trainval*"
image_directory_path: "/workspace/tao-experiments/data/training"
}
include_difficult_in_training: true
image_extension: "png"
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_fold: 0
}
以下是如果要使用序列格式的示例dataset_config元素：
dataset_config {
data_sources: {
label_directory_path: "/workspace/tao-experiments/data/training/label_2"
image_directory_path: "/workspace/tao-experiments/data/training/image_2"
}
data_sources: {
label_directory_path: "/workspace/tao-experiments/data/training/label_3"
image_directory_path: "/workspace/tao-experiments/data/training/image_3"
}
include_difficult_in_training: true
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_1"
image_directory_path: "/workspace/tao-experiments/data/val/image_1"
}
}
dataset_config中的参数定义如下：
-
data_sources：捕获要进行训练的数据集路径。如果您有多个训练数据源，可以使用多个data_sources。对于序列格式，此字段包含2个参数：
- label_directory_path：数据源标签文件夹的路径
- image_directory_path：数据源图像文件夹的路径
对于TFRecord格式，此字段包含2个参数：
- tfrecords_path：TFRecord文件的路径，这可以是匹配多个TFRecord文件的模式。
- image_directory_path：数据源图像文件夹的路径。请确保与dataset_convert命令中指定的路径一致。
-
include_difficult_in_training：指定是否在训练中包括困难的边界框。如果设置为False，将忽略困难的边界框。困难的边界框是在KITTI标签中具有非零遮挡级别的边界框。
-
image_extension：图像文件的后缀名。例如，png或jpg等。此参数仅在使用TFRecord数据集时有用。
-
target_class_mapping：此参数将标签中的类别名称映射到网络中要训练的目标类别。为每个源类别到目标类别的映射定义一个元素。此字段的目的是将相似的类别对象分组在一个类别下。例如，“car”、“van”、“heavy_truck”等可以分组在“automobile”下。其中，“key”字段是tfrecords文件中类别名称的值，“value”字段对应于网络预计要学习的值。
-
validation_data_sources：捕获要进行验证的数据集路径。如果您有多个验证数据源，可以使用多个validation_data_sources。与data_sources一样，此字段包含两个相同的参数。此参数与validation_fold互斥。
-
validation_fold：在使用TFRecord数据集进行训练时，验证数据集可以是训练数据集中的一个拆分（fold）。此参数与validation_data_sources互斥。
2.6 YOLOv3 Config
| 字段名 | 描述 | 数据类型和约束 | 推荐/典型值 |
|---|---|---|---|
| big_anchor_shape, mid_anchor_shape, and small_anchor_shape | 这些设置应该是引号内的1维数组。这些数组的元素是元组，按照宽度、高度的顺序表示预定义的锚点形状。默认情况下，YOLOv3具有九种预定义的锚点形状，分为大、中、小三组，对应不同深度的检测输出。用户应该运行kmeans命令（tao model yolo_v3 kmeans）来确定自己数据集的最佳锚点形状，并将这些锚点形状放入规范文件中。请注意，任何字段的锚点形状数量不限于3个。用户只需要在这三个字段中的每个字段中指定至少一个锚点形状。 | 字符串 | 使用tao model yolo_v3 kmeans命令生成这些形状 |
| matching_neutral_box_iou | 此字段应为0到1之间的浮点数。任何IOU值高于此浮点值的推断边界框，将在训练过程中不会回传其物体性损失。这是为了减少误报。 | 浮点数 | 0.5 |
| arch_conv_blocks | 支持的值为0、1和2。此值控制检测输出层中存在多少个卷积块。如果要复制原始DarkNet 53中的YOLOv3模型的元架构，请将此值设置为2。请注意，此配置设置仅控制YOLO元架构的大小，特征提取器的大小与此配置字段无关。 | 0、1或2 | 2 |
| loss_loc_weight, loss_neg_obj_weights, and loss_class_weights | 这些损失权重可以配置为浮点数。YOLOv3损失是定位损失、负物体性损失、正物体性损失和分类损失的总和。正物体性损失的权重设置为1，而其他损失的权重从配置文件中读取。 | 浮点数 | loss_loc_weight: 5.0 loss_neg_obj_weights: 50.0 loss_class_weights: 1.0 |
| arch | 用于特征提取的骨干。目前支持的选项有“resnet”、“vgg”、“darknet”、“googlenet”、“mobilenet_v1”、“mobilenet_v2”和“squeezenet”。 | 字符串 | resnet |
| nlayers | 特定arch中的卷积层数。对于“resnet”，支持10、18、34、50和101。对于“vgg”，支持16和19。对于“darknet”，支持19和53。所有其他网络都没有此配置，用户应删除配置文件中的此配置。 | 无符号整数 | - |
| freeze_bn | 是否在训练过程中冻结所有批处理规范化层。 | 布尔值 | False |
| freeze_blocks | 模型训练过程中要冻结的块的列表。您可以选择冻结模型中的一些CNN块，以使训练更加稳定和/或更容易收敛。块的定义对于特定体系结构是启发式的。例如，通过步进或模型中的逻辑块等方式。但是，块ID号标识了模型中块的顺序，因此在进行训练时，您不必知道块的确切位置。要牢记的一个一般原则是：块ID越小，离模型输入越近；块ID越大，离模型输出越近。您可以将整个模型分成几个块，并选择冻结其中的一个子集。请注意，对于FasterRCNN，您只能冻结在ROI池化层之前的块。任何在ROI池化层之后的层都不会被冻结。对于不同的骨干，块的数量和每个块的块ID都不同。关于如何指定每个骨干的块ID的详细说明值得一提。 | 列表（重复整数） | 见上述各种骨干的说明 |
| force_relu | 是否将所有激活函数替换为ReLU。这对于为NVDLA训练模型很有用。 | 布尔值 | False |