pytorch/accimageで高速に画像読み込み

はじめに

PyTorchの公開しているRepositoryの中にaccimageという高速に画像を読み込むためのライブラリがあります．torchvisionのDatasetクラスの中でもaccimageが使えるなら使うようになっているので，accimageをインストールすることで画像読み込みを高速化することができます．画像データを使う場合は読み込み速度はそこまで問題にならないかもしれませんが，動画を扱うときには読み込みがボトルネックになることが割とあるのでここが速くなると嬉しいんです．

意外とaccimageを使い方やどのくらい速くなるのかの情報が見つからなかったので，自分で試してみました．

accimageのインストール

accimageのRepositoryはこちら．https://github.com/pytorch/accimage
事前にIPPをインストールしておいてから，python setup.py install --userとすればいいようです．

IPPなどのインストールもanacondaで簡単にやりたかったので，forkしてRepositoryを作りました．
https://github.com/kenshohara/accimage
必要なものはipp-develとlibjpeg-turboでした．

conda install ipp-devel -c intel
conda install libjpeg-turbo -c conda-forge

あとはsetup.pyの中でanacondaでインストールしたライブラリを探しに行くように少しだけ変更しました．
anaconda_dirを手動で設定するという適当なやり方ですが，とりあえずこれでインストールできます．

# modify anaconda_dir in setup.py
python setup.py install --user

読み込み速度の比較

accimage

使用コード

%%time

% import accimage

for i in range(n_iter):
    x = accimage.Image(file_path)

結果

313x234, 15KBの画像，n_iter=1000
CPU times: user 644 ms, sys: 38.6 ms, total: 682 ms
Wall time: 686 ms

5120x2880, 502KBの画像，n_iter=10
CPU times: user 1.21 s, sys: 94.9 ms, total: 1.31 s
Wall time: 1.31 s

PIL

使用コード

%%time

% from PIL import Image

for i in range(n_iter):
    with open(file_path, 'rb') as f:
        with Image.open(f) as img:
            x = img.convert('RGB')

結果

313x234, 15KBの画像，n_iter=1000
CPU times: user 2.6 s, sys: 60.6 ms, total: 2.67 s
Wall time: 2.68 s

5120x2880, 502KBの画像，n_iter=10
CPU times: user 4.23 s, sys: 201 ms, total: 4.43 s
Wall time: 4.45 s

OpenCV

使用コード

%%time

% import cv2

for i in range(n_iter):
    x = cv2.imread(file_path)

結果

313x234, 15KBの画像，n_iter=1000
CPU times: user 2.73 s, sys: 105 ms, total: 2.83 s
Wall time: 2.85 s

5120x2880, 502KBの画像，n_iter=10
CPU times: user 4.24 s, sys: 132 ms, total: 4.37 s
Wall time: 4.39 s

比較まとめ

Wall Timeで比較

まとめ

accimage使うとやっぱりjpg読み込みが速いですね，という結果でした．動画認識だと1サンプル16フレーム，バッチサイズ128とかで学習を回したりするので，読み込みが速いくなると嬉しいです．PyTorchのDataLoader使ってマルチスレッドで読込みすればそこまで問題にはならないんですが，速いに越したことはないですね．

今回は純粋に読み込み部分の速度比較だけしましたが，データ読み込み時にはaccimageもしくはPILのImageからTensorに変換する作業も入ってきます．torchvisionのToTensorの実装を見るとPILの場合はtransposeする必要があり，そこに読み込み処理の80%の時間がかかっているようです．もしtorchvisionでaccimageを使う場合，そのtransposeをスキップすることで今回の検証以上に速くなるかもしれません．