TensorFlowの前処理を劇的に遅くするNumpy配列への変換に注意
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場合だけ桁違いに遅くなる現象が説明できません。謎ですね。
関連する情報
issueでも報告されています
Converting Tensor to Numpy array extremely slow in TF 2.0
https://github.com/tensorflow/tensorflow/issues/27692
Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー