MXNet 视频I/O读取速度提升18倍的优化策略

美图最新技术解读,不容错过

大规模视频数据的模型训练中,视频读取时间严重影响模型的训练速度。MXNet仅提供读取图像的迭代器,没有提供读取视频的迭代器,本文提出一种优化策略,可以将训练速度提升18倍。

一、前言

大规模视频数据的模型训练中,视频读取时间严重影响模型的训练速度。MXNet仅提供读取图像的迭代器,没有提供读取视频的迭代器。传统方法基于opencv或skimage直接读取原始图像,速度较慢。我们将原始图像打包成Rec格式,然后使用ImageRecordIter迭代器构建新的迭代器,具体代码实现见MTCloudVision/mxnet-videoio(https://github.com/MTCloudVision/mxnet-videoio)。使用4个Titan 1080ti GPU,优化后训练速度提升了~18倍。

MXNet框架使用迭代器器模式实现读取硬盘中图像的I/O接口。目前MXNet官方提供的读取图像的迭代器有:image.ImageIter、io.ImageRecordIter(io.ImageRecordUInt8Iter)、io.MNISTIter。MXNet的I/O接口可扩展性强,支持开发者对于图像进行打包,生成用于训练模型的迭代器。目前MXNet没有提供读取视频的I/O接口。

本文首先比较MXNet不同接口的图像I/O性能;然后在Rec图像迭代器基础上,实现视频I/O迭代器,同时对比了优化前后的性能指标。

二、图像I/O接口性能对比

MXNet三种图像I/O迭代器:

io.MNISTIter:该接口是为MNIST数据集设计的,仅支持读取MNIST图像数据,数据增强格式支持有限;

io.ImageRecordIter:支持Rec格式的数据读取。该接口同时支持多种图像增强方式。基于C++实现,执行效率较高,读取速度较快。缺点是需要将所有训练图像一次性打包成Rec格式,占用磁盘空间较大;

image.ImageIter:同时支持读取Rec和原始图像,相比以上两接口,更加灵活,同时也支持多种图像增强方式。接口基于Python实现,读取速度慢于io.ImageRecordIter接口;

我们对image.ImageIter和io.ImageRecordIter做了如下对比测试:

测试环境:

MXNet版本:0.11.0

网络结构:Inception-v3

类别(num-classes):3

GPU:titan x

测试结果:

单GPU,batchsize=128

可以看出,前两种读取方式的I\O时间主要消耗在data_iter阶段,第三种I\O时间主要消耗在update_metric阶段,且前两种时间消耗大约是第三种的1.4倍。调试ImageRecordIter接口的update_metric阶段操作,发现耗时主要集中在pred_label.asnumpy()或pred.asnumpy()操作。

多GPU(3),batchsize=128*3

可以看出,多GPU时,前两种io时间约为第三种的4.4倍。

结论:单GPU时,ImageRecordIter(Rec格式)的读取速度是其他接口的1.4倍;多GPU时,ImageRecordIter(Rec格式)是其他接口的4.4倍。原因是其他接口I/O读取数据时间是训练时间的30倍+,多GPU时,其他接口速度基本不变。如果数据集是固定的,建议使用ImageRecordIter接口进行图像读取,缺点是占用磁盘空间较大。

三、视频I/O优化性能分析

本部分介绍基于mxnet图像io迭代器ImageRecordIter的视频读取迭代器的实现方法,具体实现可以参考:MTCloudVision/mxnet-videoio(https://github.com/MTCloudVision/mxnet-videoio)。

mxnet图像I/O迭代器的输出结构:(batchsize, channel, height, width)。

我们要实现的读取视频的迭代器输出结构:(batchsize, frame_pervideo, channel, height, width),有两种方式可以实现这种迭代器,即基于opencv接口实现迭代器和对已有迭代器接口进行封装。

基于OpenCV接口实现迭代器:使用OpenCV读取视频,将读取数据进行打包成结构为(batchsize,frame_pervideo, channel, height, width)的数据。该方法优点:基于Python代码容易实现。缺点:视频读取很慢,对于大规模视频训练任务,严重影响模型的迭代效率。

封装ImageRecordIter接口:以每个视频取3帧为例,先将视频的数据封装成结构为(3*batchsize, channel, height, width)的图像数据,将标签封装成(3*batchsize,)的结构;然后调用ImageRecordIter,将图像数据reshape成(batchsize, 3, channel, height, width),并将标签进行稀疏采样成(batchsize,)的结构。

基于以上两种方法,我们做了三组性能对比实验,结果如下:

通过对比,可以看到:

基于Rec格式的数据读取速度约为使用opencv读取图像速度的18倍;

基于Rec格式的数据读取速度与GPU数正相关,4个GPU的训练速度大概是单个GPU的4倍,即多GPU训练性能提升显著;

OpenCV读取视频图像时,单GPU和多GPU的读取速度相近,即使用多GPU对训练速度的提升几乎没有帮助;

OpenCV读取视频图像,多线程(10)读取比单线程读取速度有提升,但提升有限;

以上实验结果的测试环境:

MXNet版本:1.0.1

网络结构:BN-Inception

批次数(BatchSize):50

机器:GTX1080ti

训练数据类别数(num_class):101

视频处理:视频采样3帧,每帧大小256x320

实际应用中,训练数据10W视频,每个视频截取10帧时,采用resnet-200在titan x上训练20个epoch,采用cv2.imread四个线程io需要~228小时,而基于Rec视频迭代器只需~22小时

作者简介

付志康,美图云视觉技术部门,计算机视觉工程师。

推荐一个值得参加的挑战赛

这是一个怎样的挑战赛?

美图短视频实时分类挑战赛,本次分类挑战赛是首个基于短视频数据的挑战赛,突出视频分类算法在工业界的应用,将对参赛队伍开放业界最大规模短视频数据集。

你将收获到什么?

获奖团队将获得丰厚物质奖励并由竞赛组织方颁发获奖证书:竞赛设一等奖1名、二等奖1名、三等奖2名、优胜奖5名,奖金分别为10万元、5万元、3万元、1万元;

真正参与到工业界落地应用的研发过程,对短视频研究、应用领域有更深入的了解。组委会将举办挑战赛专题研讨会,邀请获奖参赛队伍作特邀技术报告;

组委会将举办挑战赛专题研讨会,邀请获奖参赛队伍作特邀技术报告;

美图公司会对参赛获奖队伍队员开放美图Offer绿色通道

参数队伍将有机会和中国各大高校顶尖计算机视觉团队共同竞争、碰撞火花,深入了解现在视频领域业界学术发展情况

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20180529G1QVP800?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。

扫码关注云+社区

领取腾讯云代金券