TAO系列08-图像分类的模型训练与修剪

最后在训练之前,可以选择是否启动 “迁移学习” 的功能?在这个实验中也会示范这中间所得到的精度差异,让大家直接体验到迁移学习所带来的的好处。

图像分类(image classification)是视觉人工智能的最基础应用,目前TAO模型训练工具支持resnetvggmobilenet_v1mobilenet_v2googlenetcspdarknetdarknetefficientnet_b0efficientnet_b1cspdarknet_tiny10种神经网络,其中resnetvggdarknetcspdarknet还有不同结构层数的区分。

TAO启动器中只用 “classification” 这个任务指令,去面对上面所列出的十多种网络结构,其关键处就是结合“配置文件”的协助,因此这个配置文件的设置组是TAO工具的核心所在,也是本文一开始要花时间说明的部分。

cv_samples下面的classification项目,是TAO工具的图像分类模型训练的标准范例,里面提供classification.ipynb执行脚本与specs目录下两个.cfg配置文件,基本上在执行脚本里只需要修改一些路径的设定就行,其余细节都在配置文件里面进行调整,包括所选择的神经网络种类与结构。

这个项目以Pascal VOC 2012数据集作为训练源,与TAO所支持的格式并不相同,因此在训练之前还需要进行格式的转换。还好这里的转换比较简单,只要将图像分别存放在以类别所命名的目录下就可以,同时分割成训练与校验两个不同用途的数据集,在脚本里透过一段简单的Python代码就能完成,并没有什么难度。

最后在训练之前,可以选择是否启动迁移学习的功能?在这个实验中也会示范这中间所得到的精度差异,让大家直接体验到迁移学习所带来的的好处。

以下就将几个执行重点提出来说明,协助大家可以轻松地执行。

  • 配置文件内容

TAOCLI指令集只使用一个classification去面对所有的图形分类应用,其余的工作就全部交给配置文件去处理。在图像分类的配置文件里,主要有以下三个配置组:

  1. model_config:存放神经网络种类、结构(层数)与特性的内容,以下列出比较重要的部分:
  1. arch:网络名称,可使用所有已支持的网络,包括resnetvggmobilenet_v1mobilenet_v2googlenetcspdarknetdarknetefficientnet_b0efficientnet_b1cspdarknet_tiny等;
  2. n_layer:有些网络有多种结构层,例如resnet10/18/34/50/101vgg16/19darknet19/53cspdarknet19/53等;

其余参数都按照配置文件里面所设定的值,因为NGC提供的预训练模型是按照这些参数所训练的,能在迁移学习过程中得到比较好的效果。

  1. train_config:执行训练时所需要参考的
  1. train_dataset_path: 训练用数据集的位置,需要输入在容器内的完整路径
  2. val_dataset_path: 校验用数据集的位置,需要输入在容器内的完整路径
  3. pretrained_model_path: 预训练模型位置,需要输入在容器内的完整路径
  4. batch_size_per_gpu:请根据GPU显存大小进行调整
  5. n_epochs:训练的回合数
  6. n_workersCPU的并行线程数,请根据实际CPU核数量进行调整

其余参数清先安装范例所提供的设置。

  1. eval_config:评估模型时所需要用到的参数,需要输入在容器内的完整路径
  1. eval_dataset_path: 测试用数据集路径,与前面相同
  2. model_path: 在前面训练模型过程中所生成的.tlt模型文件中,挑选效果最好的一个模型来进行评估,通常最后一个的效果会最好。这里的”resnet_080.tlt”是因为训练回合数为80

其余参数清先按照范例所提供的设置。

以上关于路径部分的设置,只要遵循脚本一开始0. Set up env variables and map drives的配置规则,文件里基本上不需要做修改。其他参数的定义与设定值,请访问https://docs.nvidia.com/tao/tao-toolkit/text/image_classification.html#model-config 有完整的说明。

  • 数据集格式转换与分配

TAO支持的图像分类数据格式,是以分类名作为路径名,将所有该类的图像都放置到分类名目录之下,例如有20个分类的数据集,就会有20个分类目录。

本实验使用Pascal VOCVOCtrainval_11-May-2012.tar(1.9GB)数据集,可以用脚本提供的链接,也可以在https://pan.baidu.com/s/1JhnBCRi32xblhSSWFmF5aA(密码: gg95)下载压缩文件。

请按照脚本的数据集处理过程,包括放置路径与解压缩的指令。由于这个数据集的格式并不符合TAO的要求,因此脚本中使用一段Python代码来进行转换,详细内容请自行阅读代码,其主要功能就是将下图左边VOC路径结构转换成右边符合TAO要求的结构,存放在“formated”目录下。

TAO系列08-图像分类的模型训练与修剪

接着再将“formated”的数据,随机分割成trainvaltest三大类存放到“split”目录下,作为后面执行训练、校验与测试时所需要的数据路径。

TAO系列08-图像分类的模型训练与修剪

  • 动迁移学习功能

TAO的模型训练继承TLT的迁移学习功能,只要在配置文件中的“train_config”配置组里,指定好“pretrained_model_path:”的路径就可以,如果不想启用这个功能,只要在这个参数前面用 “#” 去关闭就行。

NGC提供300多个预训练模型,我们需要挑选合适的模型来协助执行迁移学习,这部分的下载需要使用NGC的指令来操作,在脚本中已经提供完整的安装步骤,只要执行命令块就行,然后用以下指令去下载所选的模型:

ngc registry model download-version 模型全名 --dest 本地存放路径

这里配合所使用的网络种类,挑选nvidia/tao/pretrained_classification:resnet18作为本次的预训练模型,文件大小在80MB左右。

下图左方是这个实验中,关闭迁移学习功能所训练出来的模型,经过下面评估步骤所测算出来的结果,下图右方则是在基于NGC下载的预训练模型基础上,利用迁移学习功能所训练出来的模型效果,总体来说得到的精准度提升还是很明显的。

TAO系列08-图像分类的模型训练与修剪

但如果是缺乏合适预训练模型的状况下,也不能硬是挑选不合适的模型来做迁移学习,可能得到的效果会更差。

  • 执行模型训练

当前面的准备工作都做到位之后,这个步骤就是水到渠成。执行训练所耗费的时间就只跟GPU卡的总体计算资源有关系,如果您的设备有一张以上GPU计算卡的话,可以透过TAO指令以下参数,去调整可用的计算资源:

  • --gpus  [N]:指定要调用的GPU卡数量;
  • --gpu_index 编号1 [编号2...]

例如设备上有2GPU卡时,想要全部用上的话就直接用 “--gpus 2” 参数进行调用;如果设备上有4GPU计算卡,想要指定调用0号与3号这两张卡的时,就用 “--gpus 2 --gpu_index 0 3” 参数,就能精准调用指定的GPU来进行训练。

在这个图像分类实验中,我们以NVIDIA RTX2070RTX3070进行测试,这两张计算卡都具有8GB显存。在固定回合数为80batch_size_per_gpu64的条件下,所耗费的训练时间如下:

  • GPU/3070-8G    : 5219
  • GPU/2070-8G    : 6244
  • GPU/3070+2070 : 3803
  • 修剪模型

这个步骤的目的,是要在维持足够精准度的前提下缩小模型的尺寸,越小的模型会消耗更少的计算资源,特别是内存与显存。

这对于计算资源充沛的计算设备来说,所体现的差异并不明显,但是对于计算资源相对吃紧的边缘设备来说,就有非常大的影响,甚至影响推理识别的性能,因此请自行决定是否需要执行这个步骤。

这里的要剪裁的对象是已训练好的模型,而不是从网络结构层进行融合处理,最简单的方法就是调整 “-pth(阈值),在精确度与模型大小之间取得平衡,调整的值越高就会得到越小的模型,但精准度会受到比较大的损失,这就需要经过不断尝试去调节出满足要求的模型。

以这个classification项目为例,在80回合中训练的模型大小为87.7MB,通常最后一回合训练的模型精确度会最好,因此用resnet_080.tlt模型进行修剪测试:

  • -pth=0.6:修剪后生成大小为38.5MBresnet18_nopool_bn_pruned.tlt
  • -pth=0.3:修剪后生成大小为44.2MBresnet18_nopool_bn_pruned.tlt

至于修剪后模型对精确度产生怎样的影响?就需继续往下执行再训练之后,看看评估的效果如何!

  • 模型再训练

经过修剪的模型并不能立即使用,而是作为再训练的预训练模型,再次利用迁移学习的技巧,以前面的数据集进行模型再训练的工作,唯一不一样的地方是这次使用的配置文件为classification_retrain_spec.cfg,这个配置内容与前面的classification_spec.cfg只有三个差异的地方:

  1. model_config设置组少了两个”freeze_blocks”设置;
  2. train_config设置组的“pretrained_model_path”换成前面修剪过的模型;
  3. eval_config设置组的“model_path”换成修剪后重新训练的模型

继续往下执行7. Retrain pruned models8. Testing the model两个步骤,就能得到修剪过模型的精确度评估结果。

TAO系列08-图像分类的模型训练与修剪

从上面的比较表中可以看出,这个剪裁的方式对模型文件大小的影响是很明显的,但精确度的损失并不是太大,所以这个剪裁赢视为有效果。

  • 本实验结语

这个范例主要让大家体会并熟悉TAO的配置文件与执行流程,事实上TAO所提供的范例几乎都是一样的流程,后面还有导出模型、INT8优化、生成TensorRT引擎的步骤,我们结合在后面的物件检测范例中进行说明。【完】

来源:业界供稿

0赞

好文章,需要你的鼓励

2022

04/20

16:54

分享

点赞

邮件订阅