MindSpore:【课程作业经验】基于Mindspore实现MTCNN
介绍
本项目的来源是我选修的北航与华为合作的《AI开源计算系统前沿技术》课程大作业。课程请到华为的各位专家介绍了华为目前的AI软硬件体系,并讲解了许多人工智能领域的知识。
我的大作业选题是用轻量级的网络模型backbone,实现手机端人脸检测算法,并使用MindSpore Lite在端侧推理部署。比较遗憾的是整体项目开发进度慢于预期,加上移动端目标检测APP的源码中使用了JNI等我不熟悉的接口,最后没有实现移动端的部署,仅实现了人脸检测+关键点检测神经网络的搭建、训练和测试。
前期调研之后,我决定基于MindSpore框架实现论文《Joint Face Detection and Alignment using Multi-task Cascaded Convolutional Networks》中的人脸识别和关键点检测网络MTCNN。该网络的结构并不复杂,包括PNet、RNet、ONet这三个结构相似的网络。由于我对Python和机器学习都是初学者,项目过程中遇到了不少问题。解决过程中有了一些经验心得,斗胆在此做些分享。
本项目中的MTCNN部分代码基于夜雨飘零1的Github项目和华为MindSpore官方文档撰写,包括数据集生成、网络结构定义、网络训练、模型测试的代码。
0 训练环境
笔记本训练环境:
- CPU:i7-11800H
- GPU:RTX 3070
- Windows10:
- CUDA11.6+cudnn8.x
- MindSpore1.7.0-CPU
- Pytorch1.11.0
- Ubuntu22.04:
- CUDA11.1+cudnn8.0.4
- MindSpore1.7.0-GPU
服务器训练环境:
- minsspore1.7.0-cuda10.1-py3.7-ubuntu18.04, GPU: 2*V100(64GB), CPU: 16核128GB
- minsspore1.7.0-cuda10.1-py3.7-ubuntu18.04, GPU: 1*P100(16GB), CPU: 8核64GB
1 数据集下载
WIDER Face下载训练数据压缩包WIDER Face Training Images,解压的WIDER_train文件夹放置到dataset文件夹下。并下载Face annotations,解压把里面的wider_face_train_bbx_gt.txt文件放在dataset目录下,
在Deep Convolutional Network Cascade for Facial Point Detection下载Training set并解压,将里面的lfw_5590和net_7876文件夹放置到dataset下
最终,dataset目录下应该有文件夹lfw_5590,net_7876,WIDER_train,有标注文件testImageList.txt,trainImageList.txt,wider_face_train.txt,wider_face_train_bbx_gt.txt(这四个txt文件已放在dataset文件夹下)。
2 文件夹功能说明
dataset:最开始只包含原始数据集,后续会保存用于训练每个网络的数据集。
infer_models:保存训练出的模型PNet.ckpt,RNet.ckpt,ONet.ckpt,我训练出的模型文件已经放在该文件夹下。
models:
- Loss .py:定义损失函数。
- PNet .py,RNet .py,ONet .py:定义网络。
load_models:包括PNet .py、RNet .py、ONet.py三个文件,实际上RNet.py和ONet.py与models文件夹中的文件一模一样,只有PNet做了改动。这是因为训练中PNet的输入是12*12的图片,经过卷积之后最后两维都是1,进行了squeeze操作。但在推理过程中,输入PNet的是图像金字塔,再使用squeeze会产生错误,故load_models文件夹中的PNet减少了squeeze操作。
train_PNet:包含生成训练PNet的数据集的文件generate_PNet_data.py和训练PNet的文件train_PNet.py。train_RNet和train_ONet文件夹类似。
utils:包含加载图片、处理图片的函数。
infer_camera.py:调用电脑摄像头,对摄像头得到的图片进行推理,实时显示人脸回归框和五个关键点。
infer_path.py:处理指定路径的图片,识别图片中的人脸,显示人脸回归框和五个关键点(可处理多人)。
modelToMNDIR .py:将.ckpt格式的模型转换为mindir格式。
3 模型训练过程
MTCNN是一个级联网络模型,包含PNet,RNet,ONet三个网络。这三个网络的计算结果一个比一个精确,后两个网络都会起到对其前一个网络的结果进行进一步筛选的作用。生成数据集时会调用之前网络进行推理,将之前网络的推理结果用于生成下一个网络的数据集,网络训练过程比较繁琐。
训练步骤:
- 进入train_PNet文件夹,运行generate_PNet_data.py生成训练PNet的数据集;运行train_PNet.py训练PNet。
- 进入train_RNet文件夹,运行generate_RNet_data.py生成训练RNet的数据集;运行train_RNet.py训练RNet。
- 进入train_ONet文件夹,运行generate_ONet_data.py生成训练ONet的数据集;运行train_ONet.py训练ONet。
训练结束后,如果想验证模型训练结果,可以运行infer_camera.py,调用电脑的摄像头,模型参数正确则可以显示人脸回归框和关键点。也可以运行infer_path.py,对指定图片进行推理。
4 项目心得
最开始网络训练时loss不下降,我推断出是梯度反向传播的问题,但迟迟没有解决。感谢ms技术交流群的wgx老师,帮我修改了loss函数和训练PNet的代码(训练RNet和ONet的代码是类似的),并规范了我的模型定义。
4.1 模型训练
MindSpore中数据集的加载方式和Pytorch不同,所以我自定义了数据库加载类GetDatasetGenerator来获取数据。
在Pytorch中,网络模型和损失函数是分开计算的,使用起来比较自由。MindSpore中模型的训练流程相比之下要固定很多,但也确实对代码做出了简化。因为该模型的损失函数对类别、回归框、关键点三部分的损失做加权求和,故需要自定义损失函数。在ms中自定义的损失函数要应用到网络中需要再定义一个类,我定义的这个类名为NetWithLossCell。该类中调用网络backbone得到计算结果,再将结果送到loss函数中计算损失。训练过程中如果想查看损失函数,则应该在model.train中传入callbacks,用LossMonitor来输出损失。
夜雨飘零老师的代码里训练过程中还会输出模型的精度,这在Pytorch中是直接计算的,但ms的Model类想返回两个值非常困难,所以wgx老师帮助我通过在callbacks中传入一个精度计算函数来实现了这个功能。该函数的本质是对数据集做一次eval,因为传入的是整个训练集所以非常耗时,我在训练过程中没有使用。理论上来说如果可以提取训练集的一部分传入精度计算函数就可以节约大量时间并获得精度。
4.2 API映射问题
Pytorch版代码中,MaxPool2d的ceil参数设置为True。由于mindspore中的MaxPool2d方法没有ceil这个参数,故直接使用会导致结果的shape与论文中的不同,于是我在进行pool之前先进行了padding,可以保证每层的输出结果与论文中一致。
pytorch版代码中的其他方法都可以在mindspore中找到功能甚至名称一样的方法,写法上注意一下即可。
4.3 内存溢出和推理速度过慢问题
生成训练RNet和ONet用的数据集时会采用图像金字塔的思想,这会导致输入到MindSpore中的图像shape不断变化。如果是在GRAPH_MODE下训练模型,似乎会让框架重复生成网络,导致程序的内存占用量不断上升,最终在Linux系统中导致内存耗尽程序自动退出。如果是PYNATIVE_MODE,程序的内存占用量会不断上升,但最终占用量不会像在GRAPH_MODE下那么夸张,生成RNet的数据集用了25G左右,我开了32G的虚拟内存之后就可以应付。不过生成ONet用的数据集时,因为生成数据集的写法问题,导致内存占用量太大,我在linux系统中新开的虚拟内存没办法挂载上根节点,故最后是在服务器上完成了ONet数据集的生成和ONet的训练。
关于内存溢出,还有一个我不清楚原因的问题,在minsspore1.7.0-cuda10.1-py3.7-ubuntu18.04, GPU: 1*P100(16GB), CPU: 8核64GB
这一环境上运行generate_ONet_data.py的时候发生了内存溢出的情况,64GB内存不够用,于是我开了minsspore1.7.0-cuda10.1-py3.7-ubuntu18.04, GPU: 2*V100(64GB), CPU: 16核128GB
这个环境,但在这个环境下内存最多用了48G,代码是一模一样的。可以发现,ms框架在不同的环境上运行效果有一定的区别。
最开始的训练过程中,我发现ms代码的模型处理数据集中的一张图片用时比pytorch版本代码长很多。分析后发现原因还是最开始loss不下降,意味着PNet等模型没有学到什么东西,后续生成训练RNet的数据集时就多了许多错误的回归框等数据,因此增大了数据量,增加了处理时间。后来调整模型使得模型可以正确推理,得到正确的数据集之后,ms版本代码的推理速度就正常了,并且哪怕是PYNATIVE_MODE推理速度仍然比pytorch版本代码快(因为怕推理过程中再出现内存泄露,故没敢用GRAPH_MODE训练)。