TensorFlowの前処理を劇的に遅くするNumpy配列への変換に注意

Posted On 2020-04-28

8.4k{icon} {views}

TensorFlowの前処理では、Numpy配列とTensorFlowの配列を同時に扱うことがあります。サンプル単位のTFの配列を、np.asarrayでバッチ化したら激重になってしまったので、解決策とベンチマークを調査してみました。

発生した状況

データを読み込む際に以下のようなジェネレーターを作っていました（簡略化して書いています）

環境：TF2.1.0、Numpy1.18.3

import numpy as np
import tensorflow as tf

def generator(jpeg_list, label_list):
    def gen():
        while True:
            indices = np.random.permutation(len(jpeg_list))
            for i in range(10): # 何らかのミニバッチの数
                imgs, labels = [], []
                for j in range(16): # バッチサイズ
                    idx = i * 10 + j
                    x = tf.image.decode_jpeg(jpeg_list[idx])
                    x = tf.image.resize(x, (224, 224))
                    y = np.eye(10)[label_list[idx]] # One-hot encoding
                    imgs.append(x)
                    labels.append(y)
                yield np.asarray(imgs), np.asarray(labels)
    return gen

なにがおこったか

「めちゃくちゃ重い」。1バッチ単位で1分以上かかるのでありえない。原因を探していくとTFの変数（x）をnp.asarrayしているのがボトルネックになっていました。

コードについて

KerasのGeneratorライクな書き方です。バッチ単位でyieldするのはTensorFlowの流儀には反するかもしれない。

引数のjpeg_listはJPEGファイルのバイナリを格納したlist、label_listはラベルのlistです。本来はもっとごちゃごちゃしたコードですが、簡略化のためラベルは乱数でおきます。

forループの中身がバッチ単位のサンプルの追加で、xはTFのテンソル、yがNumpy配列です。TensorFlowの関数の入力にNumpy配列は入れることは可能です（後述）。また、この逆でTensorFlowのテンソルをNumpy関数で使ったり、matplotlibに直接読み込ませることも可能です。

注意すべきなのはTensorFlowのテンソルをNumpy配列化するケースです。

TFの関数にNumpyを入れる

まずはTFとNumpyの互換性から。Numpy配列＋Numpy配列の結合をTFの関数で行うことは可能です。出力はTFのテンソルになります。

a = np.random.randn(2, 3)
b = np.random.randn(2, 3)
x = tf.concat([a, b], axis=0)
print(x)
#tf.Tensor(
#[[-0.24203849 -0.95816006 -0.27056419]
# [-1.20394309  0.97804627  0.33956476]
# [ 1.1562606  -1.1446378  -0.26468612]
# [ 0.12358135  0.29125095  0.23957379]], shape=(4, 3), dtype=float64)

普通はNumpyだけでこう書くはずです。違いはConcatしたときのデータのタイプになります。以下は出力がNumpy配列になります。

a = np.random.randn(2, 3)
b = np.random.randn(2, 3)
x = np.concatenate([a, b], axis=0)
print(x)
# [[ 1.25738828  1.32919263 -0.81560987]
#  [ 0.61928488  0.66463377 -2.99683406]
#  [ 1.17263578  0.74583488 -1.14523024]
#  [-0.5477374   2.15504587 -0.42508299]]

また、TFのテンソルとNumpy配列の結合もできます。この場合も出力はTFのテンソルになります。

a = np.random.randn(2, 3)
b = tf.random.normal(shape=(2, 3), dtype=tf.float64)
x = tf.concat([a, b], axis=0)
print(x)
# tf.Tensor(
# [[ 1.50659764  1.1113709  -1.28022306]
#  [ 1.0760203   0.84534397 -0.99987397]
#  [ 0.09307432 -0.78312794  1.04101507]
#  [ 0.48735044  0.38452375 -0.11365245]], shape=(4, 3), dtype=float64)

デフォルトのデータ型は、np.random.randnがfloat64（np.float64）、tf.random.normalがfloat32（tf.float32）であることに注意してください。

また、TFとNumpyの合成をNumpyの関数で行うこともできます。

a = np.random.randn(2, 3)
b = tf.random.normal(shape=(2, 3), dtype=tf.float64)
x = np.concatenate([a, b], axis=0)
print(x)
# [[ 1.1131192  -0.32940707  0.7312496 ]
#  [ 0.31576501  0.64231699  0.36890784]
#  [ 1.05161315  0.93724685  0.68696761]
#  [-0.60800293 -1.15293944  1.70253579]]

少し脱線しますが、TFのテンソルをmatplotlibで直接表示するのも可能です。デバッグのときに便利そうですね。

import matplotlib.pyplot as plt

x = tf.image.decode_image(tf.io.read_file("parrot.jpg"))
x = tf.image.convert_image_dtype(x, tf.float32)
print(x.shape, x.dtype) # (3648, 5472, 3) <dtype: 'float32'>

plt.imshow(x) # 直接TFのテンソル入れてOK
plt.show()

ベンチマーク

では肝心の速度はどうでしょうか？

対象：（Numpy, TF）、関数：（Numpy, TF）の4つの条件に対し、1行単位でコメントアウトしながら時間を計測してみました。公平にするためTFではGPUを切り、CPUで計測しています。224x224x3の乱数を64個結合します。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import tensorflow as tf
import time
import numpy as np

start_time = time.time()

np_data = [np.random.randn(224,224,3).astype(np.float32) for i in range(64)]
tf_data = [tf.random.normal(shape=(224, 224, 3)) for i in range(64)]

x = tf.stack(np_data, axis=0) # elapsed :  3.497114658355713
# x = np.asarray(np_data) # elapsed :  0.45649170875549316
# x = tf.stack(tf_data, axis=0) # elapsed :  0.46546053886413574
# x = np.asarray(tf_data) # elapsed :  567.6758272647858

print("elapsed : ", time.time() - start_time)

1個だけものすごい遅いのがありますね。567秒ってこれはひどい。結果をまとめます。

入力	結合	秒
Numpy	np.asarray	0.46
Numpy	tf.stack	3.50
TF	np.asarray	567.68
TF	tf.stack	0.47

入力と結合処理が同じ場合はやはり速いです。入力がNumpyでtf.stackで結合した場合も若干遅いですが、それ以上にありえないぐらい遅いのが入力がTFでnp.asarrayするケースです。

ここで気になったので、np.asarrayではなく、np.stackを使った場合を調べてみました。（各パターンごとに時間を調べました）。

x = np.stack(np_data, axis=0) # elapsed :  0.4858360290527344
x = np.stack(tf_data, axis=0) # elapsed :  0.44962191581726074

TFの変数をNumpy関数で結合しても、np.stackを使えばほとんど変わらない結果になりました。np.asarrayでもnp.stackでも出力のshapeは同じです。

ジェネレーターで比較

先程のジェネレーターの例を

np.asarrayで結合
tf.stackで結合
np.stackで結合

するケースを比較します。いずれも出力のshapeは同じです。

def generator(jpeg_list, label_list, type):
    def gen():
        while True:
            indices = np.random.permutation(len(jpeg_list))
            for i in range(10): # 何らかのミニバッチの数
                imgs, labels = [], []
                for j in range(16): # バッチサイズ
                    idx = i * 10 + j
                    x = tf.image.decode_jpeg(jpeg_list[idx])
                    x = tf.image.resize(x, (224, 224))
                    y = np.eye(10)[label_list[idx]] # One-hot encoding
                    imgs.append(x)
                    labels.append(y)
                if type == 0:
                    yield np.asarray(imgs), np.asarray(labels)
                elif type == 1:
                    yield tf.stack(imgs, axis=0), tf.stack(labels, axis=0)
                elif type == 2:
                    yield np.stack(imgs, axis=0), np.stack(labels, axis=0)
    return gen

import pickle
import time

if __name__ == "__main__":
    with open("cache/trainval_img.pkl", "rb") as fp:
        jpeg_list = list(pickle.load(fp).values()) # 何らかのJPEGを固めたもの
    labels = np.random.randint(low=0, high=10, size=(len(jpeg_list)))  # クラス（説明のため乱数）
    concat_type = 0
    dataset = tf.data.Dataset.from_generator(generator(jpeg_list, labels, concat_type),
                (tf.float32, tf.float32), ((None, 224, 224, 3), (None, 10)))

    start_time = time.time()
    for x, y in dataset:
        print(x.shape)
        print(y.shape)
        print("type :", concat_type, " = ",  time.time() - start_time)
        break

concat_typeを0, 1, 2と比較してベンチマークを行います。jpeg_listはVOC2007のtrainvalのJPEGをバイナリ化したものです。

concat_type	関数	秒
0	np.asarray	145.260
1	tf.stack	0.032
2	np.stack	0.033

やはり、np.asarrayが激重です。1回測っただけではtf.stackとnp.stackの間の差は不明でした。

ただ先程の実験を加味すると、Numpyの配列だろうがTFの変数だろうが、np.stackで結合するのがベストそうです。

結論

TFとNumpyの変数がごちゃまぜしているときは結合処理に注意が必要
np.asarrayは激重だから、TFの変数があるときは使うのは良くない
np.stackかtf.stackを使うべき。np.stackのほうがTFの変数が混じったときですら速そう（TFの変数が入力でも、tf.stackよりnp.stackのほうが速いケースがある）

「なぜnp.asarrayが重いか？」という点について。np.asarrayが新しくメモリを割り当てられているとしたら、入力がNumpyの場合も遅くなるので、入力のTF場合だけ桁違いに遅くなる現象が説明できません。謎ですね。