
作者 | 王小二@知乎(已授权)
来源 | https://zhuanlan.zhihu.com/p/397989007
编辑 | 极市平台
极市导读
本文展示迁移mxnet代码到pytorch,文章使用人脸多属性这个比较常用的方向来做演示。内容涵盖了多属性数据集如何统一,如何支持新属性的加入,基本的训练流程,简单的检测+属性演示,如何导出onnx以及caffe模型,如何比对pytorch、onnx、caffe的计算结果等内容。 >>加入极市CV技术交流群,走在计算机视觉的最前沿
本文展示迁移mxnet代码到pytorch,怎么做多属性模型。文章使用人脸多属性这个比较常用的方向来做演示。内容涵盖了多属性数据集如何统一,如何支持新属性的加入,基本的训练流程,简单的检测+属性演示,如何导出onnx以及caffe模型,如何比对pytorch、onnx、caffe的计算结果等等乱七八糟的内容。文章不是展现一些数据处理或者训练的小tricks,而是尽可能简单的展示一个基本的流程。
注意:文章展示的是我个人的训练流程,不合理之处,大家多多包涵。
0x00:人脸多属性
文章演示的人脸多属性包括了以下部分:人脸二次分数(人脸质量)、性别、年龄、微笑、戴帽子、戴眼镜、戴口罩、5点关键点,一共8个属性。示范了常见的二分类套路以及回归套路,示范了多个属性来自同一数据集或者来自不同数据集。
下面先看几张首次训练后的效果图(效果拉胯,仅供示例):
0x01:基本的训练套路
在这一节中文章会稍微详细的展示一些流程,例如如何创建数据集?多个数据集的数据怎么整合到一个训练工程里面来?一套基本完整的训练工程应该包含哪些内容?数据增强如何才能高效一点?带landmark时数据增强需要注意什么吗?训练的评测要怎么来操作?诸如此类。
那我们就从数据准备开始进行介绍。
1、数据集部分
我演示的流程采用了来自8个数据集的不同属性数据,例如从celeba中抽取了人脸质量、性别、微笑、戴帽子、戴眼镜、5点关键点属性;从imdb wiki中抽取了性别、年龄属性;从imfd&cmfd抽取了戴口罩属性。下面列举两个数据集地址,其它的就不一一列举了。
https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
https://data.vision.ee.ethz.ch/cvl/rrothe/imdb-wiki
由此可以看出,我们要使用的属性并不是每个数据集都包含了,那我们应该怎么解决多个数据集整合的问题呢?(这个问题看起来似乎很简单,但是在我第一次接触多属性训练之前,我就一直不明白标注原来还可以这么用的。我们学习检测、分类等问题时,示例代码总是会让你在同一个数据集上训练或者把相同的标签抽出来训。比如我们训一个猫狗分类器,训练的数据集就只有猫狗两个类。训练mnist时,数据集就不会有a,b,c等字母。)
实际上我们在训练多属性的时候,可以通过mask去动态选择我们需要的标签和训练目标。比如我们示例的属性就来自多个数据集,每个数据集都不包含完整的训练信息。这里我们首先讨论多个数据合并到一起的问题,后面说训练的时候怎么用mask进行选择。
我们拿到多个数据集的不同属性,第一反应可能是把当前数据集有的属性记录为真实标签,没有的数据集记录为-1或者任何一个不会出现的值即可。这样我们就会得类似:image_a:0,0,1,0 0,-1这样的标注信息。在我处理完所有数据集后,得到的标签就可以直接读入使用了。然而这个方法有一些小缺点,比如我不知道这些0,1到底是啥;比如在我处理完所有数据集后又突然要加一个新属性,这样我所有的数据集都要重新处理一遍。还有就是数据集多了之后需要占用的存储空间也会比较多。比如我们示例的图片大概200w多一点,按照这种格式存储的话,需要txt的大小就有320M这么多。
那我们就考虑有没有比较好的方法可以解决上面的小问题呢?那肯定是有的,比如我们把解析放到训练前,然后生成的数据类似image_a:age_25|gender_0等等,每个图片只包含实际的标签,并且标签使用真实意义的单词做标识。在dataloader中,我们动态去解析这些数据即可。同时我们需要的存储也减少了差不多一半左右。
还有一个问题,不同数据集的人脸裁剪区域并不相同,有的还没有box标注。这个应该怎么办呢?解决办法就是所有数据集通过统一的人脸检测器进行裁剪。示例中采用的人脸检测模型来自widerface的标注方式,检测后对box进行扩展,然后裁剪出人脸区域。
这里再提一下人脸质量这个点,示例中就类似mtcnn生成样本的方式,选择检测框和标注框iou大于某个值的做高质量样本,小于某个值的做低质量样本。这里只是做示例,不代表这样做就是可以用的(实际上还真有点效果(・。・))。
2、训练部分
这一部分包含了配置文件读取、模型构建、数据读取器、数据增强、优化器、学习率调整、损失函数构建、测试集测试、模型保存、信息打印等等。
下面逐个介绍一下。
A、配置文件读取:配置文件使用了yaml文件,配置了数据集路径、优化器信息、数据增强信息、模型构成信息等等。解析也很容易,直接open,然后yaml.safe_load加载就可以使用了,获取具体配置的时候,按照读dict的方式比如epochs = cfg['total_epochs']直接读取就行。使用yaml做简单的配置信息还是不错的,看起来比较直观。
with open(config_file, 'r') as f:
cfg = yaml.safe_load(f)
B、模型构建: 模型构建只要保证输出可以满足需求,基本的backbone可以随意选择,示例使用了分支模型,输出使用了sigmoid做二分类,线性输出做回归。基本结构如下所示:
看到这个结构,应该做个属性或者人脸对齐工作的都会想起一篇文章
https://arxiv.org/abs/1611.00851
在前司的时候,我们也搞了一个all in one 这你敢信?
C、数据读取器: 这里包含了对动态标签的解析,以及数据缓存等部分。动态标签解析部分我们需要配合之前生成的标签格式进行。按照示例的生成方式,我们首先需要对数据使用'|'进行分割,分割的第一部分就是图片地址,后面部分就是属性标签然后逐个解析属性标签即可。但是我们之前说了,标签顺序可能是动态的,所以逐个解析就不太好,于是有另外一个方法,直接搜索属性对应的值。代码实现如下:
An All-In-One Convolutional Neural Network for Face Analysis:https://arxiv.org/abs/1611.00851
C、数据读取器: 这里包含了对动态标签的解析,以及数据缓存等部分。动态标签解析部分我们需要配合之前生成的标签格式进行。按照示例的生成方式,我们首先需要对数据使用'|'进行分割,分割的第一部分就是图片地址,后面部分就是属性标签然后逐个解析属性标签即可。但是我们之前说了,标签顺序可能是动态的,所以逐个解析就不太好,于是有另外一个方法,直接搜索属性对应的值。代码实现如下:
def get_number_after_str(string, part):
number = re.findall((r"(?<=%s\:)\d+\.?\d*" % part), string) # 提取指定字符后数字
number = list(map(float, number))
if len(number) < 1:
return -1
return number[0]
这样我们只需要遍历配置文件中需要训练的属性进行解析即可。支持动态标签,标签混用,新增标签,无用标签处理。
for idx, attr_name in enumerate(self.cfg['names']):
attrs[idx] = cv_utils.get_number_after_str(attr_str, attr_name)
关于数据缓存,虽然解析标签耗时不多(我的设备上200w条数据大概1分钟可以解析完),但是有时候你希望这1分钟都不要浪费,比如你就想测一下打印日志的格式是不是合理。缓存在示例中就直接用了torch.save/load来存储dict数据,读取缓存后直接将dict信息解析到list即可无缝接入之前的环节。尽可能少的引入模块,这是我个人的习惯。是否缓存的依据目前是使用的文件修改时间做标准,不是太合理。
torch.save(cache_info, cache_path)
cache_info = torch.load(cache_path)
self.img_files = list(cache_info.keys()) # img_list
self.label_files = list(cache_info.values()) # labels_list
D、数据增强: 数据增强部分示例了基本的颜色变换、亮度调整、随机缩放、随机裁剪、水平翻转、随机旋转等等。为了充分的加速操作,部分操作被合并到一起进行了。这一点是我在yolov5的代码里白嫖过来的。需要注意的大概就是关键点的处理,需要保证变换的正确性。
随机增强后的结果
E、优化器和学习率调整: 这里示例了如何对模型中不同的数据进行操作,如卷积权重,batchnorm参数,bias参数。如何简单的增加一个学习率调整策略和warmup实现。
# Optimizer
pg0, pg1, pg2 = [], [], [] # optimizer parameter groups
for _, v in model.named_modules():
if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
pg2.append(v.bias) # biases
if isinstance(v, nn.BatchNorm2d):
pg0.append(v.weight) # no decay
elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
pg1.append(v.weight) # apply decay
if 'adam' == cfg['optimizer']:
optimizer = optim.Adam(pg0, lr=cfg['lr_base'], betas=(cfg['momentum'], 0.999))
else:
optimizer = optim.SGD(pg0, lr=cfg['lr_base'], momentum=cfg['momentum'], nesterov=True)
optimizer.add_param_group({'params': pg1, 'weight_decay': cfg['weight_decay']}) # add pg1 with weight_decay
optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
del pg0
del pg1
del pg2
# Warmup
nw = round(cfg['warmup_epochs'])
if epoch <= nw:
for _, x in enumerate(optimizer.param_groups):
x['lr'] = x['initial_lr'] * 0.1 + x['initial_lr'] * 0.9 * epoch / nw
def one_cycle(y1=0.0, y2=1.0, steps=100):
# lambda function for sinusoidal ramp from y1 to y2
return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
Warmup+OneCycle
F、损失函数构建: 损失函数基本上就是不同训练方法的灵魂了。就如检测里面ssd,yolo,anchor free等等。我们的backbone这些基本都可以选一样的。最后的区别就是如何选择标签进行训练,也就是损失函数构建。我们示例就使用交叉熵进行二分类,wingloss和SmoothL1Loss做回归。并且演示了如何解决loss出现nan;如何实现fl,qfl,vfl;如何对动态标签进行选择等。
我们在这里故意没有使用pytorch提供的BCEWithLogitsLoss函数,而是直接用了BCELoss,直接用BCELoss大家都知道在log(x)的x<=0的时候,会出现nan的情况,示例故意采用这种方式。
这个loss是给landmark回归使用的。并不适合age用,虽然看上去age也是回归问题,但是age在loss小的时候是应该更小,而不是更大。所以age的回归使用了SmoothL1Loss来进行。
这里看着和wingloss很像,其实有个比较大的区别就是在diff比较小的时候两个loss的改变并不是趋同的。
G、测试集测试: 这里包含了对测试集进行loss计算以及基本的accuracy和定位误差的计算。
测试集计算loss的话,基本按照训练集的方式跑一遍数据就行了,但是这里要注意测试的时候我们并不需要计算梯度,所以要把梯度去掉:
H、模型保存: 模型就每个epoch结束后保存一次,然后根据上面计算的accuracy等信息全局得到一个最优模型保存一下。
def save_checkpoint(model, filename):
torch.save(model.state_dict(), filename)
I、日志和打印: 演示比较粗糙包括了控制台的直接打印和tensorbord的输出,以及训练图片的显示等。如果需要的话,也可以考虑把打印同时写入文件中,示例就没做这个功能了。
至此应该就可以训练出一个演示用的pytorch模型了,下面演示一下如何对训练得到的模型进行测试,以及pytorch模型如何导出等等。
0x02:数据测试以及模型导出
1、数据测试演示了如何在图片上操作,要进行摄像头或者video测试只需要改变一下输入即可。
测试pytorch模型大概就是构建模型,load权重,选择gpu或者cpu,然后灌入数据即可。要注意的就是我们还是要用一个人脸检测器来统一输入的人脸区域。这里选择和前面生成数据的人脸检测一致。
2、在示例中演示了导出onnx,导出caffe,以及pytorch、onnx、caffe的结果比对。
导出的话pytorch都比较简单,前面和测试模型差不多,构建模型,加载权重然后按照需要设置到导出到onnx的格式就行了。
这里推荐导出后再使用onnx-simplifier进行一次操作,同时大老师还提供了一键转换各个模型数据的网页版本。这里请大声喊出来:大老师yyds!!!
https://github.com/daquexian/onnx-simplifier
https://convertmodel.com
这里不得不为这些大佬的无私点赞打call
在获得onnx模型后,我们就可以转换为caffe和进行结果比对了。转换脚本都封装好了,最终只需要像导出onnx一样,调用一个api即可。
0x03:一些后记
-
文章使用人脸属性进行示例,要迁移到其它属性应该也是ok的,比如车辆属性、行人属性等等。
-
演示代码比较简单,甚至可能有错误的地方,需要注意参考。
-
模型结构和loss函数以及各个训练参数都有很大的调整空间。
-
训练速度应该还可以更快,目前在我的设备上训练大概是1.4w imgs/s,测试大概是5w imgs/s。
-
示例代码,正在更新中...
如果觉得有用,就请分享到朋友圈吧!
公众号后台回复“CVPR21检测”获取CVPR2021目标检测论文下载~
# CV技术社群邀请函 #
备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳)
即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群
每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~

