用NAIDIA TAO工具套件打造超轻超快超准确的口罩检测模型，从训练到部署仅需1天（1）

Nvidia TLT现已升级更新为TAO（Train, Adapt, and Optimize）工具套件。

曾经听说，超过半数的机器视觉应用需求都可以通过一个目标检测模型来解决。尽管在目标检测上已经有了大量开源框架和经典模型，很多同学在进行项目落地时仍遇到了种种困难。毕竟工业应用与学术研究不同，不光要追求模型的准确性，也要兼顾性能和性价比。另外，时间成本也是很重要的。一套好的工作流可以为开发者节省大量的时间，也保证了高质量模型能快速产出和迭代。在上个月底，Nvidia发布了迁移学习工具包（Transfer Learning Toolkit）3.0版。TLT正是为开发者提供的高效模型生产工具，配合DeepStream构成了一套从训练到部署的完整解决方案。

TLT3.0的文档已汇总在下方的链接中。其中第一篇Blog就是口罩检测的案例，代码也可以在GitHub上获取。这个案例原本是为TLT2.0写的，现在我们来用TLT3.0做一遍，借此看看3.0和2.0有什么区别。

Transfer Learning Toolkit Get Starteddeveloper.nvidia.com/tlt-getting-started

相信大家对迁移学习的概念都不陌生。在机器视觉模型中，已经很少有从头训练的模型。尤其是前几年ImageNet预训练几乎成了标配。

而TLT中所包含的：

一方面是大量预训练模型，
另一方面是模型框架及训练工具。

下方表格汇总了TLT中的图像分类、目标检测、实例分割、语义分割模型框架及支持的Backbone。其中DetectNet_V2是Nvidia主推的目标检测模型，在大多数任务上都能达到业界领先的性能和准确性，适合工业应用。TLT中很多特定领域的模型都是基于DetectNet_V2的预训练模型。我们这次口罩检测也是使用DetectNet_V2来训练。

可以看到，TLT3.0中增添了一些通用的模型框架。除此之外，3.0相对2.0在接口上做了一些调整。无论是3.0还是2.0版本，TLT的主要部分都运行于Docker容器中。也就是说，TLT基本只要求主机安装显卡驱动和Docker，而cuda和cudnn都包含在容器中，不会出现什么依赖问题。因此TLT可方便部署在本地或服务器。对于nvidia-docker2的安装，很多朋友一开始会被nvidia-container-runtime，nvidia-container-toolkit等内容搞晕。在此一句话讲清楚，装nvidia-docker2就行，因为nvidia-docker2依赖其它内容，是最高层级的抽象。具体结构关系见下图。

而Docker安装Ubuntu系统仓库中的http://docker.io即可，无需安装docker-ce。这里的区别就是docker-ce是Docker官方维护的，而http://docker.io是Debian维护的，组织结构更符合Debian风格。毕竟对于Jetson平台，JetPack预装的Docker也是http://docker.io。nvidia-docker2的安装参考下面的链接即可。

Setting up NVIDIA Container Toolkitdocs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#setting-up-nvidia-container-toolkit

Installation Guide在2.0中，开发者直接在Docker容器中进行操作，而3.0则提供TLT Launcher Python包来对Docker进行操作。在实际测试中发现，这一改动至少带来两个好处，一是避免对容器内容的意外修改，二是精简了启动命令。如果用户直接运行Docker容器，往往需要输入很长的启动参数，包括需要挂载的目录，使用的显卡，分配的内存等。而通过TLT Launcher，这些配置可以预先写入~/.tlt_mounts.json配置文件。这部分更详细的信息可以参考下面两个链接：

TLT Launcherdocs.nvidia.com/metropolis/TLT/tlt-user-guide/text/tlt_launcher.html

Migrating to TLT 3.0docs.nvidia.com/metropolis/TLT/tlt-user-guide/text/migrating_v2.0_v3.0.html

使用TLT训练目标检测模型主要有三部分工作。首先是准备数据和预训练模型，然后进行TLT训练，最后进行剪枝和重训练，流程见下图。

训练口罩检测用到了4个数据集：Kaggle的医用口罩数据集、MAFA数据集、FDDB数据集、WiderFace数据集。其中前两个既包含戴口罩的人脸标注，又包含无口罩的人脸标注，主要提供了口罩类别的样本。而后两个数据集则进一步补充了无口罩人脸的样本，使得2个类别样本数量更为均衡。将数据转为KITTI格式，使用原来的Python脚本即可，这部分与TLT无关。顺便说一下，每个类别样本数量原作者只提取了6000个，再从中划分训练集和验证集。而为了充分利用数据，其实最多能各提取8256个有口罩和无口罩人脸样本作为训练集。另外提取专门的验证集。中Kaggle和FDDB数据集没有划分验证集，因此只用到MAFA和WiderFace来做验证，两个类别各能提取1121个样本。DetectNet_V2需要固定分辨率的输入，其中宽度>=480，高度>=272，且宽高均为16的倍数。在检查了数据集图像后，建议使用640*480的分辨率，如果使用16:9的比例容易使一些图像拉伸太宽。接下来将准备好的数据和标注转换为TLT需要的TFRecords，命令由2.0的tlt-dataset-convert改为3.0的tlt detectnet_v2 dataset_convert。如果提前划分了训练集和验证集，则需准备对应2个配置文件来转换，其中val_split分别为0和100。

模型的Backbone我们先用和原案例一致的ResNet18来测试，对于这种检测类别不多的简单任务，一般ResNet18就足够了。如果预先安装配置了NGC Registry CLI，预训练模型下载会方便一些。NGC CLI其实不是必需的，开发者可以前往NGC手动下载模型。

ngc registry model download-version nvidia/tlt_pretrained_detectnet_v2:resnet18 \
    --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18

下面是供大家参考的训练配置文件detectnet_v2_train_resnet18_kitti.txt，其中在数据增广中增加了一些缩放。另外TLT3.0的后处理除了原先的DBSCAN，还增加了NMS和二者结合的HYBRID模式。当然，这三种在部署到DeepStream时都是支持的。这里我们仍然使用DBSCAN，有兴趣的同学可以尝试对比其它后处理方式。

random_seed: 42
dataset_config {
  data_sources {
    tfrecords_path: "/workspace/tlt-experiments/data/tfrecords/kitti_train/*"
    image_directory_path: "/workspace/tlt-experiments/data/train"
  }
  image_extension: "jpg"
  target_class_mapping {
    key: "mask"
    value: "mask"
  }
  target_class_mapping {
    key: "no-mask"
    value: "no-mask"
  }
  validation_data_source: {
    tfrecords_path: "/workspace/tlt-experiments/data/tfrecords/kitti_val/*"
    image_directory_path: "/workspace/tlt-experiments/data/test"
  }
}


augmentation_config {
  preprocessing {
    output_image_width: 640
    output_image_height: 480
    min_bbox_width: 1.0
    min_bbox_height: 1.0
    output_image_channel: 3
  }
  spatial_augmentation {
    hflip_probability: 0.5
    zoom_min: 0.8
    zoom_max: 1.2
    translate_max_x: 8.0
    translate_max_y: 8.0
  }
  color_augmentation {
    hue_rotation_max: 25.0
    saturation_shift_max: 0.20000000298
    contrast_scale_max: 0.10000000149
    contrast_center: 0.5
  }
}

postprocessing_config {
  target_class_config {
    key: "mask"
    value {
      clustering_config {
        clustering_algorithm: DBSCAN
        dbscan_confidence_threshold: 0.9
        coverage_threshold: 0.00499999988824
        dbscan_eps: 0.20000000298
        dbscan_min_samples: 0.0500000007451
        minimum_bounding_box_height: 20
      }
    }
  }
  target_class_config {
    key: "no-mask"
    value {
      clustering_config {
        clustering_algorithm: DBSCAN
        dbscan_confidence_threshold: 0.9
        coverage_threshold: 0.00499999988824
        dbscan_eps: 0.15000000596
        dbscan_min_samples: 0.0500000007451
        minimum_bounding_box_height: 20
      }
    }
  }
}

model_config {
  pretrained_model_file: "/workspace/tlt-experiments/detectnet_v2/pretrained_resnet18/tlt_pretrained_detectnet_v2_vresnet18/resnet18.hdf5"
  num_layers: 18
  use_batch_norm: true
  objective_set {
    bbox {
      scale: 35.0
      offset: 0.5
    }
    cov {
    }
  }
  training_precision {
    backend_floatx: FLOAT32
  }
  arch: "resnet"
}

evaluation_config {
  validation_period_during_training: 10
  first_validation_epoch: 10
  minimum_detection_ground_truth_overlap {
    key: "mask"
    value: 0.5
  }
  minimum_detection_ground_truth_overlap {
    key: "no-mask"
    value: 0.5
  }
  evaluation_box_config {
    key: "mask"
    value {
      minimum_height: 20
      maximum_height: 9999
      minimum_width: 10
      maximum_width: 9999
    }
  }
  evaluation_box_config {
    key: "no-mask"
    value {
      minimum_height: 20
      maximum_height: 9999
      minimum_width: 10
      maximum_width: 9999
    }
  }
  average_precision_mode: INTEGRATE
}

cost_function_config {
  target_classes {
    name: "mask"
    class_weight: 1.0
    coverage_foreground_weight: 0.0500000007451
    objectives {
      name: "cov"
      initial_weight: 1.0
      weight_target: 1.0
    }
    objectives {
      name: "bbox"
      initial_weight: 10.0
      weight_target: 10.0
    }
  }
  target_classes {
    name: "no-mask"
    class_weight: 8.0
    coverage_foreground_weight: 0.0500000007451
    objectives {
      name: "cov"
      initial_weight: 1.0
      weight_target: 1.0
    }
    objectives {
      name: "bbox"
      initial_weight: 10.0
      weight_target: 10.0
    }
  }
  enable_autoweighting: true
  max_objective_weight: 0.999899983406
  min_objective_weight: 9.99999974738e-05
}

training_config {
  batch_size_per_gpu: 24
  num_epochs: 120
  learning_rate {
    soft_start_annealing_schedule {
      min_learning_rate: 5e-06
      max_learning_rate: 5e-04
      soft_start: 0.10000000149
      annealing: 0.699999988079
    }
  }
  regularizer {
    type: L1
    weight: 3.00000002618e-09
  }
  optimizer {
    adam {
      epsilon: 9.99999993923e-09
      beta1: 0.899999976158
      beta2: 0.999000012875
    }
  }
  cost_scaling {
    initial_exponent: 20.0
    increment: 0.005
    decrement: 1.0
  }
  checkpoint_interval: 10
}

bbox_rasterizer_config {
  target_class_config {
    key: "mask"
    value {
      cov_center_x: 0.5
      cov_center_y: 0.5
      cov_radius_x: 0.40000000596
      cov_radius_y: 0.40000000596
      bbox_min_radius: 1.0
    }
  }
  target_class_config {
    key: "no-mask"
    value {
      cov_center_x: 0.5
      cov_center_y: 0.5
      cov_radius_x: 1.0
      cov_radius_y: 1.0
      bbox_min_radius: 1.0
    }
  }
  deadzone_radius: 0.400000154972
}

训练命令由2.0的tlt-train detectnet_v2改为3.0的tlt detectnet_v2 train，如果使用Turing或Ampere架构的显卡，建议加上--use_amp参数来启用自动混合精度训练。这样可以利用Tensor Cores来加速训练，同时占用更少的显存（意味着有需要时可以增大Batch Size）。在我笔记本的移动版3080上，开启混合精度训练后性能翻倍，训练总共耗时2小时40分。如果一切顺利的话，应该可以得到类似下图的结果，mAP大约0.87。

其实这时我们已经有了一个可以使用的模型，而工业应用中，我们往往希望模型在不损失精度的情况下跑的更快，消耗算力更小，因为这意味着更高的性价比。为此主要的做法有剪枝和量化，而TLT对此都是支持的。做剪枝的命令由2.0的tlt-prune改为3.0的tlt detectnet_v2 prune。pth取0.1时得到结果如下。

可以看到，剪枝后的模型大小仅为原模型的18%。根据以前关于模型剪枝的研究，一般模型剪掉80%再经过重新训练都能恢复到接近原模型的精度。在口罩检测原案例中，对剪枝后的模型做了重训练和INT8 Calibration。而INT8 Calibration其实会稍微损失一点精度。如果已经决定了要在部署时使用INT8推理，这里更好的做法是使用量化感知训练（QAT）。这样不仅可以使转换到INT8后最大限度接近FP32模型的精度，还能省去最后的Calibration操作。使用QAT做重训练的参考配置文件detectnet_v2_retrain_resnet18_kitti_qat.txt如下。

random_seed: 42
dataset_config {
  data_sources {
    tfrecords_path: "/workspace/tlt-experiments/data/tfrecords/kitti_train/*"
    image_directory_path: "/workspace/tlt-experiments/data/train"
  }
  image_extension: "jpg"
  target_class_mapping {
    key: "mask"
    value: "mask"
  }
  target_class_mapping {
    key: "no-mask"
    value: "no-mask"
  }
  validation_data_source: {
    tfrecords_path: "/workspace/tlt-experiments/data/tfrecords/kitti_val/*"
    image_directory_path: "/workspace/tlt-experiments/data/test"
  }
}


augmentation_config {
  preprocessing {
    output_image_width: 640
    output_image_height: 480
    min_bbox_width: 1.0
    min_bbox_height: 1.0
    output_image_channel: 3
  }
  spatial_augmentation {
    hflip_probability: 0.5
    zoom_min: 0.8
    zoom_max: 1.2
    translate_max_x: 8.0
    translate_max_y: 8.0
  }
  color_augmentation {
    hue_rotation_max: 25.0
    saturation_shift_max: 0.20000000298
    contrast_scale_max: 0.10000000149
    contrast_center: 0.5
  }
}

postprocessing_config {
  target_class_config {
    key: "mask"
    value {
      clustering_config {
        clustering_algorithm: DBSCAN
        dbscan_confidence_threshold: 0.9
        coverage_threshold: 0.00499999988824
        dbscan_eps: 0.20000000298
        dbscan_min_samples: 0.0500000007451
        minimum_bounding_box_height: 20
      }
    }
  }
  target_class_config {
    key: "no-mask"
    value {
      clustering_config {
        clustering_algorithm: DBSCAN
        dbscan_confidence_threshold: 0.9
        coverage_threshold: 0.00499999988824
        dbscan_eps: 0.15000000596
        dbscan_min_samples: 0.0500000007451
        minimum_bounding_box_height: 20
      }
    }
  }
}

model_config {
  pretrained_model_file: "/workspace/tlt-experiments/detectnet_v2/experiment_dir_pruned/resnet18_nopool_bn_detectnet_v2_pruned.tlt"
  num_layers: 18
  use_batch_norm: true
  load_graph: true
  objective_set {
    bbox {
      scale: 35.0
      offset: 0.5
    }
    cov {
    }
  }
  training_precision {
    backend_floatx: FLOAT32
  }
  arch: "resnet"
}

evaluation_config {
  validation_period_during_training: 10
  first_validation_epoch: 10
  minimum_detection_ground_truth_overlap {
    key: "mask"
    value: 0.5
  }
  minimum_detection_ground_truth_overlap {
    key: "no-mask"
    value: 0.5
  }
  evaluation_box_config {
    key: "mask"
    value {
      minimum_height: 20
      maximum_height: 9999
      minimum_width: 10
      maximum_width: 9999
    }
  }
  evaluation_box_config {
    key: "no-mask"
    value {
      minimum_height: 20
      maximum_height: 9999
      minimum_width: 10
      maximum_width: 9999
    }
  }
  average_precision_mode: INTEGRATE
}

cost_function_config {
  target_classes {
    name: "mask"
    class_weight: 1.0
    coverage_foreground_weight: 0.0500000007451
    objectives {
      name: "cov"
      initial_weight: 1.0
      weight_target: 1.0
    }
    objectives {
      name: "bbox"
      initial_weight: 10.0
      weight_target: 10.0
    }
  }
  target_classes {
    name: "no-mask"
    class_weight: 8.0
    coverage_foreground_weight: 0.0500000007451
    objectives {
      name: "cov"
      initial_weight: 1.0
      weight_target: 1.0
    }
    objectives {
      name: "bbox"
      initial_weight: 10.0
      weight_target: 10.0
    }
  }
  enable_autoweighting: true
  max_objective_weight: 0.999899983406
  min_objective_weight: 9.99999974738e-05
}

training_config {
  batch_size_per_gpu: 24
  num_epochs: 120
  # enabling qat when retraining the model.
  enable_qat: true
  learning_rate {
    soft_start_annealing_schedule {
      min_learning_rate: 5e-06
      max_learning_rate: 5e-04
      soft_start: 0.10000000149
      annealing: 0.699999988079
    }
  }
  regularizer {
    type: L1
    weight: 3.00000002618e-09
  }
  optimizer {
    adam {
      epsilon: 9.99999993923e-09
      beta1: 0.899999976158
      beta2: 0.999000012875
    }
  }
  cost_scaling {
    initial_exponent: 20.0
    increment: 0.005
    decrement: 1.0
  }
  checkpoint_interval: 10
}

bbox_rasterizer_config {
  target_class_config {
    key: "mask"
    value {
      cov_center_x: 0.5
      cov_center_y: 0.5
      cov_radius_x: 0.40000000596
      cov_radius_y: 0.40000000596
      bbox_min_radius: 1.0
    }
  }
  target_class_config {
    key: "no-mask"
    value {
      cov_center_x: 0.5
      cov_center_y: 0.5
      cov_radius_x: 1.0
      cov_radius_y: 1.0
      bbox_min_radius: 1.0
    }
  }
  deadzone_radius: 0.400000154972
}

重训练得到的结果如下，mAP约为0.865，与剪枝前的模型接近。

最后使用tlt detectnet_v2 export命令就可以导出模型和量化文件以供DeepStream部署。关于使用DeepStream对视频流实时分析，可以阅读本专栏之前的文章。

2GB版Jetson Nano也能多路视频实时分析？用DeepStream榨干史上最低价Jetson209 赞同 · 12 评论文章正在上传…重新上传取消

最终得到的模型大小不到10MB，性能嘛，还没测试。。。因为我之前用TLT2.0时已经做出了更小的模型，大小仅为1.2MB，在Jetson Xavier NX上可以达到300+FPS，效果如下。

仅1.2MB大小的口罩检测289 播放 · 1 赞同视频正在上传…重新上传取消

这里最大的改动是把Backbone换成了MobileNet_V2，其它操作类似。有兴趣的同学可以试试用TLT3.0能不能达到同样的效果。MobileNet_V2中加入了残差结构，一般来说，原模型Backbone为ResNet的改成MobileNet_V2都有较高的成功率。而原模型Backbone为VGG这种级联结构的做轻量化改造换成MobileNet_V1也容易一次性成功。

可以看到，在整套流程中，其实不需要开发者写任何代码。开发者的主要任务是准备数据和各种配置文件。TLT3.0还进一步丰富了针对特定领域的预训练模型，其中目标检测类模型几乎都是基于DetectNet_V2。表格中是这些模型在不同硬件上的推理性能。

表格上半部分是TLT2.0中就提供的模型，部分已在之前的专栏文章中展示过。而3.0中又增加了行人分割、车牌检测和识别、人脸关键点检测、目光检测、心率检测、手势识别、情绪识别等模型。其中预训练的车牌检测和识别是支持中国车牌的，有兴趣的同学可以试试看。

TLT3.0作为一次主版本号更新，不仅完善了机器视觉部分，还增加了对NLP的支持。Nvidia的多模态对话AI应用框架Jarvis目前已开放公测。开发者可以使用Nemo或TLT为Jarvis训练模型。关于TLT3.0的NLP部分之后有机会再为大家介绍。

用TAO工具套件打造超轻超快超准确的口罩检测模型，从训练到部署仅需1天 - 知乎

用NAIDIA TAO工具套件打造超轻超快超准确的口罩检测模型，从训练到部署仅需1天（1）

猜你喜欢