人工智能/深度学习

利用真实和合成数据准备目标检测模型和 NVIDIA 迁移学习工具箱

漫长而繁琐的数据采购工作一直在阻碍人工智能的创新,特别是在计算机视觉领域,因为计算机视觉依赖于标记的图像和视频进行训练。但是现在你可以通过使用人工智能快速生成合成数据来启动你的机器学习过程。

有了 AI . revire 综合数据平台 ,你可以在找到和标注正确的真实照片所需时间的一小部分时间内创建你需要的精确训练数据。在 AI . Revire 的照片级真实感 3D 环境中,您可以为所有可能的场景生成数据,包括难以到达的地方、不寻常的环境条件以及罕见或独特的事件。

训练数据生成 包括标签。选择所需的类型,例如二维或三维边界框、深度遮罩等。测试模型后,可以返回到平台以快速生成其他数据以提高准确性。以快速、反复的循环进行测试和重复。

我们想在 NVIDIA 迁移学习工具包 3 . 0 中测试人工智能的性能。最初,我们着手复制研究论文 RarePlanes :合成数据飞行 中的结果,该论文使用合成图像创建目标检测模型。我们在 TLT 中发现了新的工具,使得创建更轻量级的模型成为可能,这些模型与原始论文中的模型一样精确,但速度要快得多。

在这篇文章中,我们将向您展示如何使用 TLT 量化感知训练和模型修剪来实现这一点,以及如何自己复制结果。我们将向您展示如何创建飞机探测器,但您应该能够针对自己的各种卫星探测场景对模型进行微调。

This synthetic image shows annotations with some bounding boxes on several aircrafts
图 1 .一个合成图像,其特征是注释,表示飞机类型、机翼形状和其他区别特征。

访问卫星探测模型

要复制这些结果,您可以克隆 GitHub 存储库,并使用附带的 Jupyter 笔记本。

克隆以下 repo :

git clone git@github.com:aireveries/rareplanes-tlt.git ~/Code/rareplanes-tlt 

创建 conda 环境:

conda create -f env.yaml 

激活模型:

source activate rareplanes-tlt 

启动 Jupyter :

jupyter notebook 

学习目标

  • 使用 AI . revire 平台生成合成数据并与 NVIDIA TLT 一起使用。
  • 使用合成数据训练高精度模型。
  • 使用 TLT 优化推理模型。

先决条件

我们使用 python3 . 8 . 8 测试了代码,使用 anaconda4 . 9 . 2 管理依赖项和虚拟环境。这些代码可能适用于不同版本的 Python 和其他虚拟环境解决方案,但我们尚未测试这些配置。我们使用了 ubuntu18 . 04 . 5lts 和 NVIDIA 驱动程序 460 . 32 . 03 和 CUDA 版本 11 . 2 。 TLT 需要驱动程序 455 . xx 或更高版本。

  • 设置 NVIDIA 容器工具箱/ NVIDIA -docker2 。有关更多信息,请参阅 NVIDIA 容器工具包安装指南
  • 设置 NGC 以下载 NVIDIA Docker 容器。遵循 TLT 用户指南 中的步骤 4 和 5 。有关 NGC CLI 工具的更多信息,请参阅 CLI 安装
  • 至少有 250 GB 的可用硬盘空间来存储数据集和模型权重。

下载数据集

有关 RarePlanes 数据集内容的更多信息,请参阅 RarePlanes 公共用户指南

对于本教程,您只需要下载数据的一个子集。下面的代码示例将在 Jupyter 笔记本 中执行。首先,创建文件夹:

!mkdir -p data/real/tarballs/{train,test}
 !mkdir -p data/synthetic 

现在使用此函数从 Amazon S3 下载数据集,提取它们,并验证:

def download(s3_path, out_folder, out_file_count):
     rel_file_path = Path('data') / Path(s3_path.replace('s3://rareplanes-public/', ''))
     rel_folder = rel_file_path.parent / out_folder
     num_files = !ls $rel_folder | wc -l
     try:
         if int(num_files[0]) == out_file_count:
             print(f'{s3_path} already downloaded and extracted')
         else:
             raise Exception
     except:
         if not rel_file_path.exists():
             print('Starting download')
             !aws s3 cp $s3_path $rel_file_path;   
         else:
             print(f'{s3_path} already downloaded')
         print('Extracting...')
         !cd {rel_folder.parent}; pv {rel_file_path.name} | tar xz;
         print('Removing compressed file.')
         !rm $rel_file_path 

然后下载数据集:

download('s3://rareplanes-public/real/tarballs/metadata_annotations.tar.gz', 
          'metadata_annotations', 9)
 download('s3://rareplanes-public/real/tarballs/train/RarePlanes_train_PS-RGB_tiled.tar.gz', 
          'PS-RGB_tiled', 11630)
 download('s3://rareplanes-public/real/tarballs/test/RarePlanes_test_PS-RGB_tiled.tar.gz', 
           'PS-RGB_tiled', 5420)
 !aws s3 cp --recursive s3://rareplanes-public/synthetic/ data/synthetic 

从 COCO 转换成 KITTI 格式

TLT 使用 KITTI 格式进行目标检测模型训练。 RarePlanes 是 COCO 格式的,因此必须从 Jupyter 笔记本中运行转换脚本。这将转换真实列车/测试和合成列车/测试数据集。

%run convert_coco_to_kitti.py 

data/kitti 中的每个数据集都应该有一个文件夹,其中包含 KITTI 格式的注释文本文件和指向原始图像的符号链接。

设置 TLT 安装

笔记本有一个生成 ~/.tlt_mounts.json 文件的脚本。有关各种设置的详细信息,请参阅 运行启动器

{
     "Mounts": [
         {
             "source": "/home/patrick.rodriguez/Code/rareplanes-tlt",
             "destination": "/workspace/tlt-experiments"
         }
     ],
     "Envs": [
         {
             "variable": "CUDA_VISIBLE_DEVICES",
             "value": "0"
         }
     ],
     "DockerOptions": {
         "shm_size": "16G",
         "ulimits": {
             "memlock": -1,
             "stack": 67108864
         },
         "user": "1001:1001"
     }
 } 

将数据集处理为 TFR 记录

必须将 KITTI 标签转换为 TLT 使用的 TFRecord 格式。笔记本中的 convert_split 函数可帮助您批量转换所有数据集:

 def convert_split(name):
     !tlt detectnet_v2 dataset_convert --gpu_index 0 \
         -d /workspace/tlt-experiments/specs/detectnet_v2_tfrecords_{name}.txt \
         -o /workspace/tlt-experiments/data/tfrecords/{name}/{name}
 You can then run the conversions:
 convert_split('kitti_real_train')
 convert_split('kitti_real_test')
 convert_split('kitti_synthetic_train')
 convert_split('kitti_synthetic_test') 

下载 ResNet18 卷积主干

使用 NGC 帐户和命令行工具,您现在可以下载模型:

下载 ResNet18 卷积主干

使用 NGC 帐户和命令行工具,您现在可以下载模型:

!ngc registry model download-version nvidia/tlt_pretrained_detectnet_v2:resnet18 

模型现在位于以下路径:

./tlt_pretrained_detectnet_v2_vresnet18/resnet18.hdf5

使用真实数据进行基准测试

以下命令将启动训练并将结果记录到可跟踪的文件中:

!tlt detectnet_v2 train --key tlt --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_real.txt \
     -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_real_amp16 \
     -n resnet18_real_amp16 \
     --use_amp > out_resnet18_real_amp16.log 

执行以下命令:

tail -f ./out_resnet18_real_amp16.log 

培训完成后,您可以使用笔记本中定义的函数获取模型的相关统计信息:

 get_model_param_counts('./out_resnet18_real_amp16.log')
 best_epoch = get_best_epoch('./out_resnet18_real_amp16.log')
 best_epoch 

您可以得到如下输出:

 Total params: 11,197,893
 Trainable params: 11,188,165
 Non-trainable params: 9,728
 Best epoch and map50 metric: (79, 94.2296) 

要在测试集或其他数据集上重新评估经过训练的模型,请运行以下操作:

!tlt detectnet_v2 evaluate --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_evaluate_real.txt \
     -m /workspace/tlt-experiments/{best_checkpoint} \
     -k tlt 

输出应该如下所示:

 Validation cost: 0.001133
 Mean average_precision (in %): 94.2563
  
 class name      average precision (in %)
 ------------  --------------------------
 aircraft                         94.2563
  
 Median Inference Time: 0.003877
 2021-04-06 05:47:00,323 [INFO] __main__: Evaluation complete.
 Time taken to run __main__:main: 0:00:27.031500.
 2021-04-06 05:47:02,466 [INFO] tlt.components.docker_handler.docker_handler: Stopping container. 

用合成数据做实验

  !tlt detectnet_v2 train --key tlt --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth.txt \
     -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_amp16 \
     -n resnet18_synth_amp16 \
     --use_amp > out_resnet18_synth_amp16.log 

您可以通过运行: !cat out_resnet18_synth_amp16.log | grep -i aircraft 来查看每个历元的结果

输出示例:

 aircraft                         58.1444
 aircraft                         65.1423
 aircraft                         64.3203
 aircraft                         68.1934
 aircraft                         71.5754
 aircraft                         68.5568 

用实际数据对综合训练模型进行微调

现在,用 10% 的真实数据来微调最佳性能的合成数据训练模型。为此,必须首先创建 10% 分割。

 %run ./create_train_split.py
 convert_split('kitti_real_train_10') 

然后使用此函数将模板规范中的检查点替换为仅合成训练中性能最佳的模型。

 with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10.txt', 'r') as f_in:
     with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_replaced.txt', 'w') as f_out:
         out = f_in.read().replace('REPLACE', best_checkpoint)
         f_out.write(out) 

你现在可以开始 TLT 训练了。在上一节中,仅对合成数据进行训练的模型的最佳性能历元开始微调。

!tlt detectnet_v2 train --key tlt --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_replaced.txt \
     -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_finetune_10_amp16 \
     -n resnet18_synth_finetune_10_amp16 \
     --use_amp > out_resnet18_synth_finetune_10_amp16.log 

训练完成后,您将看到 91-93%mAP50 之间的最佳 epoch ,这将使您接近仅真实模型的性能,只有 10% 的真实数据。

在笔记本中,有一个命令用于评估测试集上性能最佳的模型检查点:

 !tlt detectnet_v2 evaluate --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_evaluate_real.txt \
     -m /workspace/tlt-experiments/{best_checkpoint} \
     -k tlt 

您应该看到如下输出:

2021-04-06 18:05:28,342 [INFO] iva.detectnet_v2.evaluation.evaluation: step 330 / 339, 0.05s/step
 Matching predictions to ground truth, class 1/1.: 100%|█| 14719/14719 [00:00<00:00, 15814.87it/s]
  
 Validation cost: 0.001368
 Mean average_precision (in %): 91.8094
  
 class name      average precision (in %)
 ------------  --------------------------
 aircraft                         91.8094
  
 Median Inference Time: 0.004137
 2021-04-06 18:05:30,327 [INFO] __main__: Evaluation complete.
 Time taken to run __main__:main: 0:00:30.677440.
 2021-04-06 18:05:32,469 [INFO] tlt.components.docker_handler.docker_handler: Stopping container. 
This chart has three bars, including “real data”, “synthetic only”, and “synthetic + 10% real”.
图 2 .对合成+ 10% 真实数据的训练与对 100% 真实数据的训练结果几乎一致。

数据增强 正在对人工智能的模型训练进行微调。遐想的合成数据仅占原始真实数据集的 10% 。如您所见,这种技术生成的模型与仅对真实数据进行训练的模型一样精确。这意味着在真实的标记数据上大约节省了 90% 的成本,并且使您不必忍受冗长的标记和 QA 过程。

修剪模型

训练了一个性能良好的模型之后,现在可以减少权重的数量,以减少文件大小和推理时间。 TLT 包括 易于使用的修剪工具

其中一个值得玩味的参数是 -pth ,它设置了神经元修剪的阈值。设置得越高,修剪的参数就越多,但在某一点之后,精度指标可能会降得太低。我们发现 0 . 5 的值对这些实验有效,但在其他数据集上可能会发现不同的结果。

 !mkdir -p detectnet_v2_outputs/pruned
  
 !tlt detectnet_v2 prune \
     -m /workspace/tlt-experiments/{best_checkpoint} \
     -o /workspace/tlt-experiments/detectnet_v2_outputs/pruned/pruned-model.tlt \
     -eq union \
     -pth 0.5 \
     -k tlt 

现在可以计算修剪后的模型:

!tlt detectnet_v2 evaluate --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_evaluate_real.txt \
     -m /workspace/tlt-experiments/detectnet_v2_outputs/pruned/pruned-model.tlt \
     -k tlt > out_pruned.txt 

现在您可以看到还有多少参数:

get_model_param_counts('./out_pruned.txt') 

您应该看到如下输出:

 Total params: 3,372,973
 Trainable params: 3,366,573
 Non-trainable params: 6,400 

这是 70% ,比原来的模型,其中有 1120 万个参数小!当然,删除这么多参数会降低性能,您可以验证:

 !cat out_pruned.txt | grep -i aircraft
  
 aircraft                         68.8865 

幸运的是,您可以通过重新训练修剪后的模型来恢复几乎所有的性能。

再培训模特

与前面一样,有一个模板规范来运行此实验,它只需要您填写修剪模型的位置:

 with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain.txt', 'r') as f_in:
     with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_replaced.txt', 'w') as f_out:
         out = f_in.read().replace('REPLACE', 'detectnet_v2_outputs/pruned/pruned-model.tlt')
         f_out.write(out) 

现在可以重新训练修剪后的模型:

 !tlt detectnet_v2 train --key tlt --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_replaced.txt \
     -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_finetune_10_pruned_retrain_amp16 \
     -n resnet18_synth_finetune_10_pruned_retrain_amp16 \
     --use_amp > out_resnet18_synth_finetune_10_pruned_retrain_amp16.log 

在本实验的运行中,表现最好的 epoch 达到 91 . 925map50 ,与原来的非运行实验基本相同。

2021-04-06 19:33:39,360 [INFO] iva.detectnet_v2.evaluation.evaluation: step 330 / 339, 0.05s/step
 Matching predictions to ground truth, class 1/1.: 100%|█| 17403/17403 [00:01<00:00, 15748.62it/s]
 Validation cost: 0.001442
 Mean average_precision (in %): 91.9849
  
 class name      average precision (in %)
 ------------  --------------------------
 aircraft                         91.9849
  
 Median Inference Time: 0.003635
 2021-04-06 19:33:41,479 [INFO] __main__: Evaluation complete.
 Time taken to run __main__:main: 0:00:31.869671.
 2021-04-06 19:33:43,607 [INFO] tlt.components.docker_handler.docker_handler: Stopping container. 

量化模型

这个过程的最后一步是量化修剪后的模型,这样您就可以使用 TensorRT 获得更高级别的推理速度。我们提供量化感知培训( QAT )规范模板:

with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat.txt', 'r') as f_in:
     with open('./specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat_replaced.txt', 'w') as f_out:
         out = f_in.read().replace('REPLACE', 'detectnet_v2_outputs/pruned/pruned-model.tlt')
         f_out.write(out) 

进行 QAT 培训:

!tlt detectnet_v2 train --key tlt --gpu_index 0 \
     -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat_replaced.txt \
     -r /workspace/tlt-experiments/detectnet_v2_outputs/resnet18_synth_finetune_10_pruned_retrain_qat_amp16 \
     -n resnet18_synth_finetune_10_pruned_retrain_qat_amp16 \
     --use_amp > out_resnet18_synth_finetune_10_pruned_retrain_qat_amp16.log 

使用 TLT 导出工具导出为 INT8 量化 TensorRT 格式:

!tlt detectnet_v2 export \
   -m /workspace/tlt-experiments/{best_checkpoint} \
   -o /workspace/tlt-experiments/detectnet_v2_outputs/qat/resnet18_detector_qat.etlt \
   -k tlt  \
   --data_type int8 \
   --batch_size 64 \
   --max_batch_size 64\
   --engine_file /workspace/tlt-experiments/detectnet_v2_outputs/qat/resnet18_detector_qat.trt.int8 \
   --cal_cache_file /workspace/tlt-experiments/detectnet_v2_outputs/qat/calibration_qat.bin \
   --verbose 

此时,您可以使用 TensorRT 评估量化模型:

!tlt detectnet_v2 evaluate -e /workspace/tlt-experiments/specs/detectnet_v2_train_resnet18_kitti_synth_finetune_10_pruned_retrain_qat_replaced.txt \
                            -m /workspace/tlt-experiments/detectnet_v2_outputs/qat/resnet18_detector_qat.trt.int8 \
                            -f tensorrt 

查看输出:

2021-04-06 23:08:28,471 [INFO] iva.detectnet_v2.evaluation.tensorrt_evaluator: step 330 / 339, 0.33s/step
 Matching predictions to ground truth, class 1/1.: 100%|█| 21973/21973 [00:01<00:00, 16161.54it/s]
  
 Validation cost: 0.549463
 Mean average_precision (in %): 91.5516
  
 class name      average precision (in %)
 ------------  --------------------------
 aircraft                         91.5516
  
 Median Inference Time: 0.000840
 2021-04-06 23:08:33,182 [INFO] __main__: Evaluation complete.
 Time taken to run __main__:main: 0:02:13.453132.
 2021-04-06 23:08:34,768 [INFO] tlt.components.docker_handler.docker_handler: Stopping container. 

结论

我们对这些结果印象深刻。 AI . reviree 的 综合数据平台 只有真实数据集的 10% ,使我们能够获得与在完整真实数据集上训练时相同的性能。这意味着大约节省了 90% 的成本,更不用说在采购上节省的时间了。现在生成所需的合成数据需要几天,而不是几个月。

TLT 还使参数计数减少了 25 . 2 倍,文件大小减少了 33 . 6 倍,性能( QPS )提高了 174 . 7 倍,同时保留了 95% 的原始性能。 TLT 的功能对于修剪和量化特别有价值。

转到 AI.Reverie ,为您的项目下载 综合训练数据 ,然后开始 TLT 培训。

Tags