在NVIDIA DGX Station上利用TLT训练口罩识别模型

关于西浦国际创新港

西交利物浦国际创新港致力于构建创新创业环境，使西交利物浦大学成为联结中国与世界创新创业的全球无缝衔接驱动中心。

国际创新港已经创建了一个基于大学师生创业想法的早期孵化器，并正在加速营造开放创新实验室。国际创新港将打造创新创业基础设施，以服务于大学和社区，以及包括西交利物浦大学校友和外部企业的创新需求。

西交利物浦国际创新港与江苏省政府“创建苏南国家创新示范区核心区”行动计划紧密连接。凭借其在苏州工业园区的战略定位，西交利物浦国际创新港具有参与创建苏南国家创新示范区核心区行动计划的区位优势，并获得来自当地、江苏省、国家创新创业资金的支持。

关于NVIDIA DGX Station

NVIDIA DGX Station是专为办公室场景设计的深度学习超级计算机，其性能是市面上功能最强大的工作站的两倍。

DGX Station采用液冷设计，内部包含四块Tesla V100 GPU、NVLink互联技术、20480个NVIDIA CUDA core，这一切共同造就了500 TFLOPS的超强深度学习能力。

风冷工作站的噪音行业标准大约为45分贝。DGX Station完胜这一标准，达到35分贝，也就是一般办公室通风系统的噪音量。

为什么我们选择用DGX Station？

当我们尝试复现NVIDIA在Github上提供的”戴口罩识别“开源项目（网址：https://github.com/NVIDIA-AI-IOT/face-mask-detection）时，使用NVIDIA的TLT迁移学习模型训练工具的过程，执行到“! tlt-train detectnet_v2”训练指令时，出现”/usr/local/bin/tlt-train: line 32: 577Illegal instruction (core dumped)tlt-train-g1 ${PYTHON_ARGS[*]}“的错误。

经过几番搜索与咨询NVIDIA工程师之后，确认detectnet_v2使用Tensorflow为后台骨干，需要支持AVX2指令集的CPU上才能正常执行，而我们训练设备上的两颗Intel Xeon E5-2620v2 CPU是不支持这个指令集，因此无法顺利完成detectnet_v2的训练，必须寻找其他资源。

所幸在西交利物浦大学国际创新港，有台NVIDIA高阶DGX Station，配置2颗IntelE5-2698v4 CPU、256GB DDR4内存，最重要的是装载有4片32GB显存的NVIDIA V100计算卡，并且配置NVLINK数据交换加速汇流排，将模型训练性能提高到极致。

十分感谢西交利物浦大学科研生产力和创新办公室的协助，为我们提供这么好的计算资源，让这次的“口罩识别”项目复现得以顺利进行。首先看一下这台设备的GPU配置状况，直接执行“nvidia-smi”就能看到有4片32GB显存的TeslaV100-DGXS版本（如下图），总共有20,480个CUDA计算核与2560个Tensor计算核，计算密度应该是这个区域里面最领先的。

DGX工作站里面配置的NVLINK汇流排能将数据传输性能在往上提升一个级别，现在就执行“nvidia-smi topo -m”指令，看看这台DGX工作站的NVLINK拓扑结构，显示如下图，可以看到任意两张V100之间都有NVLINK直接连接。

需要弄清楚的工作流程：

这次口罩识别数据集有1122张图像数据，在640图像尺寸执行120周期（epoch）训练，只花了8分钟左右的时间就完成，同样的训练时间在装有单片RTX2070/8G计算卡上，大约话费4个小时时间，这中间就相差大约30倍。

因为训练时间大幅度缩短，让我们能在很短时间内完成项目的数据训练工作，包括模型剪裁与再训练的任务。

接下来我们就在这台DGX上，根据开源项目的引导来进行口罩识别项目的执行，整个项目内容主要切分为以下6大部分：

l 下载4组数据集（dataset）并进行整理

l 下载TLT镜像、建立tlt容器与启动Jupyter服务

l 在Jupyter界面执行数据集转换成KITTI结构与tfrecords格式

l 从NGC下载预训练模型

l 在tlt容器中执行模型训练与优化

l 将模型部署到Jetson设备的DeepStream执行推理识别

事实上这几个步骤里面，困难度最大的是以下4部分：

1.数据集下载、整理与转换，这部分在项目引导的内容并没有交代清楚，我们经过一番摸索与实验之后，确认这部分的流程。

2.TLT镜像下载、建立容器、启动里面Jupyter等过程，比较繁琐。

3.项目执行过程的众多路径设计，由于TLT标准范例中设置的路径都很长，最多会超过100个字符，会令众多初学者迷失。

4. 项目的设定文件名称都很长，相似度也很高，文件里面也有多个路径需要修改，也存在路径名设置问题。

因此本文主要针对上述4个问题，提供使用经验的分享。至于最后面部署（复制）到目标设备Jetson Nano上，转换成DeepStream能调用的TensorRT加速引擎，然后进行口罩识别的推理计算，是不在DGX工作站上操作，会另有专文介绍。

现在就开始在DGX工作中上执行口罩识别的模型训练任务。

数据集（dataset）下载、整理：在容器外处理

事实上这个环节是整个项目中最繁琐一环，因为项目引导中只提供4个数据库链接网址，却没有明确地指出需要在每个网页中下载哪些文件？而这部分的下载又不能直接使用“wget”下载指令，因此不能在Jupyter里面操作，最好在容器外去下载、解压缩、整理之后，才能得到符合项目要求的目录结构。

以下是我们的执行步骤，至于一些目录的设定，请自行调整适合您环境的结构，最终的数据集相关结构，请参考https://github.com/NVIDIA-AI-IOT/face-mask-detection下面的data_utils/data_tree.txt。

1. 在DGX上创建~/tlt_work目录，这个目录在后面映射到容器内的/workspace/host去，为容器内外建立起数据的桥梁。

2. 在~/tlt_work目录下建立dataset目录，以存放过程中所有需要的数据集内容

3. 在dataset目录下建立Kaggle、MAFA、FDDB、WiderFace这4个目录，存放项目提供的四个下载链接的数据集内容。

(1) Kaggle：从这里下载数据之前，需要先申请账号，进入项目提供链接之后，（https://www.kaggle.com/ivandanilovich/medical-masks-dataset-images-tfrecords）点选右上角“Download”按钮后下载一个238MB的压缩文件，存放在~/tlt_work/dataset/Kaggle里面，解压缩之后所需要的是“Medical Mask Dataset”下面的images与labels两个目录的数据，将这两个目录移到Kaggle下面，变成~/tlt_works/dataset/Kaggle/images与~/tlt_works/dataset/Kaggle/labels

(2) MAFA：这个数据集压缩文件存放在百度网盘上，链接与提取码为https://pan.baidu.com/s/1qiflOq65yMNNTXyt8020GA提取码4fz6，点击输入提取码就能下载，解压缩之后里面还有四个压缩文件，将这四个压缩文件全部解压缩，会生成train与test各自的images与labels共四个目录，存放在~/tlt_works/dataset/MAFA

(3) FDDB：点击链接（http://vis-www.cs.umass.edu/fddb）后下载网页中“Download the database”下面的“Orginal, unannotated set of images”与“Faceannotations”两个压缩文件，存放在~/tlt_works/dataset/FDDB目录下然后家解压缩，让这个目录下的路径结构为2002、2003、FDDBfolds这三个目录。

(4) WiderFace：点击连接（http://shuoyang1213.me/WIDERFACE/）之后，到“Download”下方下载“WIDER FaceTraining Images”与“WIDER FaceValidation Images”两个图片集压缩文件，还有下面的“Face annotations”标注文件的压缩文件，总共三个压缩文件，全部解压缩之后放到~/tlt_works/dataset/WiderFace目录下。

将以上数据结构与进行比对与调整，确认无误之后就可以进到下一部分，将数据图像与标注文件抽取成为KITTI格式，然后转成tfrecords格式，因为这些转换可以在容器内进行，因此下个步骤就先建立TLT容器。

下载TLT镜像与建立tlt容器：

这个步骤里面还有几个小步骤需要执行：

1. 自行在工作站上安装好Docker与nv-docker容器管理工具

2. 个人请先到NGC上申请一个账号，后面调用NGC资源时需要用到

3. 执行容器下载指令：因为本项目使用TLT 2镜像，因此执行以下指令：

br

4.以TLT2的镜像建立<容器名>，这里我们给定容器名为tlt2

$ docker run -it --name=tlt2 --net=host --gpus=all \
  -v "~/tlt_work":"/workspace/host" \
 nvcr.io/nvidia/tlt-streamanalytics:v2.0_py3 /bin/bash

这样就能直接建立容器，并且启动容器与进入容器。

注："~/tlt_work"是本机上容器外的实体路径，"/workspace/host"是容器内映射的路径。

5.进入容器内启动Jupyter Notebook服务，请执行以下指令：

$ ngc registry resource download-version \ "nvidia/gtcfallngcdemo:1.0.0"
$ jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root

6.在容器外启动浏览器，输入“127.0.0.1:8888”就能开启Jupyter界面，请自行处理密码设定的部分。

开启Jupyter环境后，容器以/workspace作为启动的根目录，因此我们需要将工作位置切换到步骤4与容器外产生对应的“host”目录下。

7.下载开源内容以建立项目

在Jupyter浏览器右上角点选“New”新增一个“Terminal”终端，执行以下指令：

$ cd host
$ git clone https://github.com/NVIDIA-AI-IOT/face-mask-detection.git

就会看到以下的信息，表示项目已经下载。

此时可以在容器外的本机上tlt_work下面看到新增一个“face-mask-detection”目录，后面的项目都会在这个目录下操作。

8.启动项目的Notebook

进入“host/face-mask-detection”，打开里面的face-mask-detection.ipynb，然后按照提示逐步进行接下去的工作，这个过程中有几处需要进行修改的地方，我们会一步一步根据我们设定的路径去说明，您请调整成适合您的路径。

9.首先在“0. Set up env variables”下方指令格中设定路径参数，请根据您实际的路径进行设定：

• %env KEY=”您在NGC上获取的一组秘钥”，请自行处理
• %env USER_EXPERIMENT_DIR=/workspace/host/face-mask-detection 说明：实际上就是这个项目的工作目录，后面过程中会在这下面建立pre-trained、unpruned、pruned、final四组模型存放的目录
• %env DATA_DOWNLOAD_DIR= /workspace/host/face-mask-detection/data 说明：这个路径是我们存放数据集的路径，下面包括前面的四组数据集、过程中会生成的KITTI格式数据集，以及最终需要的tfrecords数据格式
• %env SPECS_DIR= /workspace/host/face-mask-detection/tlt_specs 说明：这里存放训练过程中各阶段所需要的配置文件，在这个项目中的tlt_specs目录下有6个配置文件，在训练过程中需要修改部分内容的路径设定。
• %env NUM_GPUS=4 说明：系统可使用的GPU数量，现在我们使用的这台DGX工作站上有4片V100计算卡，可以全部都用上

设定好这些环境变量之后，就可以按“Shift-Enter”执行这道指令。

将数据集与转换成KITTI格式与TFRecords格式

1. 在face-mask-detection.ipynb添加2道执行指令，

(1) 安装该项目的依赖库，请添加一个“代码块”，执行以下指令：

!pip3 install -r requirements.txt

(2) 因为是在容器里面，因此需要将前面4组数据集的路径转换成容器内的对应路径，其内容如下：

!python3 data2kitti.py \
--kaggle-dataset-path /workspace/host/dataset/Kaggle \
--mafa-dataset-path /workspace/host/dataset/MAFA \
--fddb-dataset-path /workspace/host/dataset/FDDB \
--widerface-dataset-path /workspace/host/dataset/WiderFace \
--kitti-base-path $USER_EXPERIMENT_DIR/KITTI \
--category-limit 2 \
--tlt-input-dims_width 960 \
--tlt-input-dims_height 544 \
--train

然后会在face-mask-detection目录下生成KITTI结构的数据集。

2.接下去就要将KITTI格式转换成 tf records 格式，作为训练用途。

(1) 执行“!cat $SPECS_DIR/detectnet_v2_tfrecords_kitti_trainval.txt”指令后，会显示这个配置文件的内容，首先看到里面的路径显示

root_directory_path: "/home/data/train"

与我们的设置不一样，所以需要修改，可以直接从Jupyter的文件列表中进入tlt_specs目录，打开下面的detectnet_v2_tfrecords_kitti_trainval.txt就可以直接编辑，将路径修改为

root_directory_path: "/workspace/host/face-mask-detection/KITTI/train"

存档完就能执行下一步，将

(2) 执行“!tlt-dataset-convert”指令后，会将转换成tfrecords格式的数据，存放在face-mask-detection/data/tfrecords/kitti_trainval

(3) 执行“!ls -rlt$DATA_DOWNLOAD_DIR/tfrecords/kitti_trainval/”指令，就会显示转换的输出是否成功。

从NGC下载预训练模型

这个步骤相对简单，只要您在NGC里申请的秘钥正确，并且下载的路径确认之后，执行应该就不会出现问题，因此不多作解释。

在tlt容器中执行模型训练与优化

这部分是整个计算量最大的部分，也是考验设备性能的环节，比较繁琐的部分是每个阶段都使用不同的配置文件，里面都有需要修改的路径，下面简单列出每个环境的配置文件与需要修改的地方，只要这些地方都处理好之后，剩下就是计算所需要的时间问题了。

以下所有配置文件都存放在face-mask-detection/tlt_specs目录下，设定给环境变量“$SPECS_DIR”，请到这个目录去获取配置文件。文件里的路径必须提供容器内的绝对路径，

1. 训练与校验：detectnet_v2_train_resnet18_kitti.txt，第一次训练时使用，需要修改部分有以下几处：

(1) tfrecords_path:填入前面执行tfrecords转换后的路径，注意最后面的“*”号。

"/workspace/host/face-mask-detection/data/tfrecords/kitti_trainval/*"

(2) image_directory_path: 作为训练用的KITTI格式图像数据

"/workspace/host/face-mask-detection/KITTI/train"

(3) pretrained_model_file:从NGC下载的与训练模型

"/workspace/host/face-mask-detection/pretrained_resnet18/tlt_pretrained_detectnet_v2_vresnet18/resnet18.hdf5"

2. 重新训练与校验：detectnet_v2_retrain_resnet18_kitti.txt，因为后面需要经过剪裁步骤，剪裁完需要重新训练，所使用的配置文件与前一个大致相同，有些细节参数可以进行调整，主要修改路径如下：

(1) tfrecords_path: 与detectnet_v2_train_resnet18_kitti.txt相同

(2) image_directory_path: 与detectnet_v2_train_resnet18_kitti.txt相同

(3) pretrained_model_file:这部分就要修改成前面未经剪裁（unpruned）所训练的模型路径，内容如下：

"/workspace/host/face-mask-detection/experiment_dir_unpruned/weights/resnet18_detector.tlt"

其他的步骤基本上就按照face-mask-detection.ipynb的流程执行就可以，这里就不多做说明。