numpy.memmapを用いたメモリ節約法

kerasにはImageDataGeneratorという、データオーグメンテーションのための便利なクラスが用意されています。flowメソッドに画像データを渡しておけば、ジェネレータを呼び出す度に加工された状態の画像のバッチを返してくれるというものですが、メモリに乗り切らないほどの巨大なデータセットを扱う場合、画像データを用意する箇所が障害となりこの仕組みが利用できません。代替手段としてflow_from_dataframeというメソッドも用意されているのですが、実際にこれを利用するとエポック毎に画像のデコードを行うため非常に時間がかかります。

本記事ではnumpyのメモリマッピング機能を利用して、高速かつ省メモリを実現しつつ、データオーグメンテーションを可能にする方法を紹介します。

numpy.memmap

メモリマッピングとは、物理的にディスク上に存在するファイルを仮想記憶空間にマッピングする技術であり、ファイル全体をメモリに読み込まずに小さなセグメントにアクセスするために使用されます。numpyにおいてはndarrayのサブクラスでもあり、ndarrayを期待されるほとんどの場面で利用することができます。以下、具体的な使い方を説明します。

まず、確保するディスク領域を予め指定しておく必要があります。

import numpy as np
import cv2

X_train, X_test, y_train, y_test = train_test_split(df["filename"],
                                                   df[target],
                                                   test_size=0.2,
                                                   random_state=0)

X_train_shape = (len(X_train), h, w, c)
X_test_shape = (len(X_test), h, w, c)

train_mmap_path = "./train.mmap"
test_mmap_path = "./test.mmap"

memmapオブジェクトはndarrayと同様に扱うことができるので、ndarrayへの代入と同じ書き方でファイルへ書き出します。

memmap = np.memmap(train_mmap_path, dtype="uint8", mode="w+", shape=(X_train_shape))
for i, filename in enumerate(X_train):
    img = cv2.imread(filename)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, dsize=(h, w))
    memmap[i][:] = img

memmap = np.memmap(test_mmap_path, dtype="uint8", mode="w+", shape=(X_test_shape))
for i, filename in enumerate(X_test):
    img = cv2.imread(filename)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, dsize=(h, w))
    memmap[i][:] = img

事前にこのように記述したスクリプトで保存したら、学習時には次のように読み込みます。書き出しと読み込みを分けておけば、画像の並びに変更がない限りファイルを再利用できます。

X_train = np.memmap(train_mmap_path, dtype="uint8", mode="r", shape=(X_train_shape))
X_test = np.memmap(test_mmap_path, dtype="uint8", mode="r", shape=(X_test_shape))

あとはこれをflowメソッドに渡せば良いだけに思えますが、実際にはうまくいきません。flowメソッドはNumpyArrayIteratorオブジェクトを返すのですが、そのコンストラクタ内に
self.x = np.asarray(x)
の一行があるために、やはりデータセット全体をメモリに乗せようとしてしまいます。

ImageDataGeneratorの修正

次のようなIteratorを定義します。複雑に見えますが、基本的には先に述べたオリジナルのNumpyArrayIteratorの問題点を取り除いただけのものになります。

from keras_preprocessing.image.iterator import Iterator
from tensorflow.keras.utils import Sequence

class MemmapIterator(Iterator, Sequence):
    def __init__(self,
                 x,
                 y,
                 image_data_generator,
                 batch_size=32,
                 shuffle=False,
                 seed=None,
                 dtype='float32'):
        self.dtype = dtype
        x_misc = []

        self.x = x
        if y is not None:
            self.y = np.asarray(y)
        else:
            self.y = None
        self.image_data_generator = image_data_generator
        super().__init__(x.shape[0],
                         batch_size,
                         shuffle,
                         seed)

    def _get_batches_of_transformed_samples(self, index_array):
        batch_x = np.zeros(tuple([len(index_array)] + list(self.x.shape)[1:]),
                           dtype=self.dtype)
        for i, j in enumerate(index_array):
            x = self.x[j]
            params = self.image_data_generator.get_random_transform(x.shape)
            x = self.image_data_generator.apply_transform(
                x.astype(self.dtype), params)
            x = self.image_data_generator.standardize(x)
            batch_x[i] = x

        output = (batch_x,)
        if self.y is None:
            return output[0]
        output += (self.y[index_array],)
        return output

これをImageDataGeneratorを継承したクラス内で定義したメソッドで呼び出します。

from tensorflow.keras.preprocessing.image import ImageDataGenerator

class MyImageDataGenerator(ImageDataGenerator):
    def flow_from_memmap(self,
                         x,
                         y=None,
                         batch_size=32,
                         shuffle=True):
        return MemmapIterator(
            x,
            y,
            self,
            batch_size=batch_size,
            shuffle=shuffle)

あとはflowの代わりに新しく作成したflow_from_memmapメソッドを使用するだけで、通常のImageDataGeneratorと同様に使用することができます。もちろんデータオーグメンテーションも可能です。

datagen = MyImageDataGenerator()
generator = datagen.flow_from_memmap(X_train, y_train)

valid_datagen = MyImageDataGenerator()
valid_generator = valid_datagen.flow_from_memmap(X_test, y_test, shuffle=False)

history = model.fit(generator,
                    validation_data=valid_generator,
                    epochs=epochs)

パフォーマンスの比較

今回、自作のジェネレータの性能を計測するために64x64サイズの約70万枚のPNG画像を用意しました。データ量としてはPNGの状態で4GB以上、uint8なら8GB以上です。

	flow_from_dataframe	flow_from_memmap
memmap書き出し（初回）	-	12550s
画像チェック	76s	-
1epoch目	6858s	2485s
2epoch目	3116s	826s
3epoch目以降	907s	456s

flow_from_memmapはflow_from_dataframeと比べ、2倍から3倍程度高速となりました。今回は比較のため画像データを事前に整形していますが、元の画像データが入力サイズよりも大きい場合、差は大きくなります。

おわりに

今回はNumpyのメモリマッピング機能を使えば、画像を一々読み出すよりも高速にできるという記事でした。ただし、表を見て貰えば分かる通り書き出し部分が読み出し以上に時間がかかるため、データを再利用することが前提となります。そのためKFoldのように訓練データの中身を入れ替えるような手法は使えませんし、モデルの入力サイズが変更された場合も再度書き出しの必要があります。

また、上の表では2epoch目以降が1epoch目よりも速いという結果が示されていますが、これはディスクキャッシュを利用している可能性があり環境によってはそれほど高速化されないかもしれません。

numpy.memmap自体はテーブルデータなどにも応用が効きます（こちらの記事のやり方は画像以外にも利用できます）ので、ディスク容量に余裕がある方は試してみてはいかがでしょうか。