Pillow-SIMDを使った画像処理やディープラーニングの高速化の紹介

Pillow-SIMDを使ってマルチプロセス化なしで画像処理を高速化する方法を紹介します。PyTorchが対応している高速化方法ですが、実はKerasでも動かせました。簡単ながら痒いところに手が届くのでおすすめです。

きっかけ

Pillow（PIL）を使って画像処理や前処理をすることが多いのですが、Python特有の悩みとして

デフォルトだとCPUに対してマルチコアの処理をしてくれない

というのがあります。特にディープラーニングでは、JPEGの読み込みが速度のボトルネックとなることが多いので、画像の読み込みや前処理の高速化ができるととても嬉しいのです。

PillowやOpenCVをmultiprocessingで並列化して高速化するという方法もありますが、マルチプロセス化はデバッグが面倒になるので、ライブラリ一個で完結してくれるとすごいありがたいのです。

そんなときにおすすめなのが、Pillow-SIMDというライブラリです。
https://github.com/uploadcare/pillow-simd

PyTorchの情報探したときに見つけたものです。中身はPILフォークなので、純粋にPillowの代替品として使えますし、前処理でPillowを使っているKerasでも使えたりします（末尾のColabノートブックを参考にしてください）。

インストール方法

たった2行でOK。従来のPillowをPillow-SIMDで置き換えます。

pip uninstall -y pillow
pip install pillow-simd

Windows10の場合はインストールで失敗することがあるので、非公式のバイナリを落としてきます関連Issue

pip install 保存ディレクトリ\Pillow_SIMD-5.3.0.post0-cp37-cp37m-win_amd64.whl

とすればOKです。

実験

Cats and Dogsデータセット

犬猫画像を25000枚ほど集めた「Cats and Dogsデータセット」を実験台として用います。画像分類の初歩として有名なデータセット。中身は全てJPEG画像です。

通常のPillowの場合

「LANCZOS法でリサイズ→画像のシャープ化」という若干重めの処理を行います。

import time
import glob
from PIL import Image, ImageEnhance
from tqdm import tqdm

def normal_pillow():
    files = sorted(glob.glob("PetImages/Cat/*.jpg"))
    files += sorted(glob.glob("PetImages/Dog/*.jpg"))
    start_time = time.time()
    list = []
    for f in tqdm(files):
        try:
            with Image.open(f) as img:
                img = img.resize((128, 128), Image.LANCZOS)
                img = ImageEnhance.Sharpness(img).enhance(2.0)
        except:
            list.append(f)
            continue
    print("経過時間 = ", time.time()-start_time, "s")
    return list

ColabのGPU環境で計測したところ（PillowはGPUによる高速化はありません）、処理時間は139秒となりました。

またKerasで浅めのニューラルネットワークを訓練します。犬と猫の画像を分類します。（前後に壊れたファイルの削除を挟んでいます）

from tensorflow.keras import layers
import tensorflow.keras as keras

def create_shallow_network():
    def conv_bn_relu(x, ch):
        x = layers.Conv2D(ch, 5, padding="same")(x)
        x = layers.BatchNormalization()(x)
        x = layers.Activation("relu")(x)
        return x
    
    input = layers.Input((128, 128, 3))
    x = conv_bn_relu(input, 16)
    x = layers.AveragePooling2D(4)(x)
    x = conv_bn_relu(x, 32)
    x = layers.AveragePooling2D(4)(x)
    x = conv_bn_relu(x, 64)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(2, activation="softmax")(x)
    return keras.models.Model(input, x)

def train():
    model = create_shallow_network()
    model.compile("adam", "categorical_crossentropy", ["acc"])
    gen = keras.preprocessing.image.ImageDataGenerator(rescale=1.0/255
             ).flow_from_directory("PetImages", target_size=(128, 128), batch_size=128)
    model.fit_generator(gen, steps_per_epoch=25000//128, epochs=5)    

1エポックあたり77秒かかっています。これはJPEGの読み込みがボトルネックになっているせいで、Pillow-SIMDを使えばもっと高速化できます。

Epoch 4/5
195/195 [==============================] - 77s 395ms/step - loss: 0.4427 - acc: 0.7932
Epoch 5/5
195/195 [==============================] - 77s 396ms/step - loss: 0.4169 - acc: 0.8099

Pillow-SIMDの場合

Pillow-SIMDをインストールし、ランタイムを再起動します。

pip uninstall -y pillow
pip install pillow-simd

先程と同じ、「読み込み→リサイズ→シャープ化」の処理をします。コードは全く変更する必要はありません。Pillow自体を置き換えているためです。

import time
import glob
from PIL import Image, ImageEnhance
from tqdm import tqdm

def simd_pillow():
    files = sorted(glob.glob("PetImages/Cat/*.jpg"))
    files += sorted(glob.glob("PetImages/Dog/*.jpg"))
    start_time = time.time()
    for f in tqdm(files):
        try:
            with Image.open(f) as img:
                img = img.resize((128, 128), Image.LANCZOS)
                img = ImageEnhance.Sharpness(img).enhance(2.0)
        except:
            continue
    print("経過時間 = ", time.time()-start_time, "s")    

結果は61秒となりました。139秒から2倍以上の高速化です。ちなみにColabのCPUは2コアなのを考えると、それっぽい結果になります。

KerasのコードもそのままでPillow-SIMDに対応してくれます。

Epoch 4/5
195/195 [==============================] - 53s 273ms/step - loss: 0.4747 - acc: 0.7764
Epoch 5/5
195/195 [==============================] - 54s 277ms/step - loss: 0.4439 - acc: 0.7936

1エポックあたり77秒から53秒に短縮することができました。KerasのImageDataGeneratorが内部的にPillowを使っているので、Pillow-SIMDに置き換えてあげることで結果的にCPUを効率よく使ってくれるようになります（公式サポートではないので、もしかしたらうまくいかない処理があるかもしれません）。

前処理が重くなったり、CPUのコア数・スレッド数が多くなればもっと高速化の効果は出やすいかと思われます。

※ただし、必ずしもCPU100％まで使ってくれることではないようです。CPUの多い環境で計測したところ、通常のPillow→Pillow-SIMDへの変更で、CPU使用率が10%から20%まで上がりました。そのため並列化が意味がなくなったというわけではないと思います。このコードがforループを使ったシーケンシャルな書き方なので、そのせいもあるかもしれません。

もっと速くする

Kerasの場合はfit_generatorにuse_multiprocessing（スレッドベースのマルチプロセス化：デフォルトでFalse）、workers（スレッドベースのワーカー数：デフォルトで1）という引数があるので、環境によってはこの値を変えることでさらに高速化（CPUをもっと使ってくれる）ようになります。

Colab環境では、use_multiprocessing=Trueとすると逆に遅くなってしまったため、workers=4と増やしてみました。

def train_multiworker():
    model = create_shallow_network()
    model.compile("adam", "categorical_crossentropy", ["acc"])
    gen = keras.preprocessing.image.ImageDataGenerator(rescale=1.0/255
             ).flow_from_directory("PetImages", target_size=(128, 128), batch_size=128)
    model.fit_generator(gen, steps_per_epoch=25000//128, epochs=5, workers=4) #ワーカーの数を4にした

その結果、1エポック当たり42秒まで短縮することができました。

このように、本気で速度出したい場合は、マルチプロセス化、ワーカー数の増加などとセットで使うといいかもしれません。Pillow-SIMD単体でも十分効果はあるでしょう。

まとめ

[秒]	通常のPillow	Pillow-SIMD
リサイズ＋シャープ化	139	61
CNNの1エポック	77	53
+ worker=4	-	42

Pillow-SIMD、速くて簡単。とても便利！　以上！

コード

今回用いたColabのノートブックはこちらにあります。

お知らせ

技術書典6で頒布したモザイク本の通販を下記URLで行っています。会場にこられなかったけど欲しいという方は、ぜひご利用ください。

『DeepCreamPyで学ぶモザイク除去』通販
https://note.mu/koshian2/n/naa60d5c9ebba

ディープラーニングや機械学習における画像処理の基本や応用を学びながら、モザイク除去技術DeepCreamPyを使いこなし、自分で実装するまでを目指す解説書です。