大规模视频数据的模型训练中，视频读取时间严重影响模型的训练速度。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帧为例，先将视频的数据封装成结构为(3batchsize, channel, height, width)的图像数据，将标签封装成(3batchsize,)的结构；然后调用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小时。

作者介绍：付志康，美图云视觉技术部门，计算机视觉工程师。

本文转载自美图技术公众号。

原文链接：https://mp.weixin.qq.com/s/Nq-fZY1L_ULO5DtBVg8eAw

创作场景

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