MXNet 視頻I/O讀取速度優化

大規模視頻數據的模型訓練中，視頻讀取時間嚴重影響模型的訓練速度。MXNet僅提供讀取圖像的迭代器，沒有提供讀取視頻的迭代器。傳統方法基於opencv或skimage直接讀取原始圖像，速度較慢。我們將原始圖像打包成Rec格式，然後使用ImageRecordIter迭代器構建新的迭代器，具體代碼實現見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迭代器，同時對比了優化前後的性能指標。

1. 圖像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

可以看出，前兩種讀取方式的IO時間主要消耗在data_iter階段，第三種IO時間主要消耗在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介面進行圖像讀取，缺點是佔用磁碟空間較大。

2. 視頻I/O優化性能分析

本部分介紹基於mxnet圖像io迭代器ImageRecordIter的視頻讀取迭代器的實現方法，具體實現可以參考：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小時。

本文作者：付志康，計算機視覺工程師，美圖雲視覺技術部門，轉載請註明出處。