Tana Gone
Tana Gone
2 min read

Categories

TF1スタイルのコードを高速化するにはTF2(あるいはTF1.x最終版1.15)で採用された入力パイプラインを採用すれば良さそうだ

tf.data.Dataset使ってみる

Denseへの入力が4、Weightsによって行列積Mによって変換された出力が1であるNNを訓練する。 訓練データがN個あればdatasetは次のようになる。 WはYの初期値を得るために乱数を使って作る。

> N = 10; D = 4; M = 1
> X = tf.random.normal((N, D))
> X
> <tf.Tensor: shape=(10, 4), dtype=float32, numpy=
> array([[ 1.7476182 ,  0.76557606,  1.1793045 ,  0.6604082 ],
>   [-1.1182905 ,  1.2941122 , -0.45307   , -0.12836482],
>   [ 1.3309119 ,  1.3667184 ,  2.4617343 , -0.44162795],
>   [ 0.09327319,  0.90526474,  0.6416567 , -0.17305161],
>   [-1.4737343 , -2.0090747 , -0.1392107 ,  0.32674456],
>   [ 0.25767186, -0.19082204, -1.9285886 , -0.47859055],
>   [-0.67001045,  0.02660223, -0.6753374 ,  0.04300369],
>   [-0.26570737,  1.748417  , -0.41708243, -0.4584445 ],
>   [ 0.8148399 ,  0.22168487,  0.01016708, -0.8258183 ],
>   [-0.34755087,  0.21631955,  0.36211625, -1.142048  ]],
>  dtype=float32)>
> W = tf.random.normal((D, M))
> Y = tf.matmul(X, W)
> # ここまでが準備
> # Batch Size = 3とするとBatch #1, 2, 3, 4が出来る。最後の#4はBatch Size = 1
> ds_mini = tf.data.Dataset.from_tensor_slices((X, Y)).shuffle(N).batch(3)\
.prefetch(tf.data.AUTOTUNE)
> for x_batch, y_batch in ds_full:
> ...    print(x_batch)
> ...    print(y_batch)

tf.Tensor(
[[ 0.25767186 -0.19082204 -1.9285886  -0.47859055]
 [ 0.09327319  0.90526474  0.6416567  -0.17305161]
 [ 0.8148399   0.22168487  0.01016708 -0.8258183 ]], shape=(3, 4), dtype=float32)
tf.Tensor(
[[ 0.3514151 ]
 [-0.2899059 ]
 [ 0.49854147]], shape=(3, 1), dtype=float32)
 ...以下略

nn.learn(ds_mini)

> model = tf.keras.Sequential([
> ...  tf.keras.layers.Flatten(),
> ...  tf.keras.layers.Dense(10)
> ...])
> model.compile(optimizer='adam',
> ... loss='mean_squared_error',
> ... metrics=['mse'])
> model.fit(ds_mini, epochs=1)
> 4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 1ms/step - loss: 2.2483 - mse: 2.2483
> <keras.src.callbacks.history.History object at 0x11e6c0e10>
> model.fit(ds_mini, epochs=2)
> Epoch 1/2
> 4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 749us/step - loss: 2.2319 - mse: 2.2319
> Epoch 2/2
> 4/4 ━━━━━━━━━━━━━━━━━━━━ 0s 640us/step - loss: 2.2156 - mse: 2.2156
> <keras.src.callbacks.history.History object at 0x11e688050>

入力パイプラインはCPUとGPUを交互動作ではなく同時動作させる

TF1のコードスタイルですか?それともTF2ですか?それともTF1, TF2共通のコードスタイルですか?

tf.data: TensorFlow 入力パイプラインの構築 TensorFlow Core

train, test = tf.keras.datasets.fashion_mnist.load_data()

images, labels = train
images = images/255.0
labels = labels.astype(np.int32)
fmnist_train_ds = tf.data.Dataset.from_tensor_slices((images, labels))
fmnist_train_ds = fmnist_train_ds.shuffle(5000).batch(32)

model = tf.keras.Sequential([
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(10)
])

model.compile(optimizer='adam',
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), 
  metrics=['accuracy'])

model.fit(fmnist_train_ds, epochs=2)

理由(主なポイント)

  • tf.keras.datasets.fashion_mnist.load_data() でデータを読み込んでいる → TF2の標準的なKeras API。
  • tf.data.Dataset.from_tensor_slices() + .shuffle() + .batch() でデータパイプラインを作成している → TF2で強く推奨される効率的な入力パイプライン(TF1ではほとんど使われず、TF2以降の本流)。
  • tf.keras.Sequential([…]) でモデル定義。
  • model.compile() + model.fit(dataset, epochs=…) で訓練 → これがTF2の標準的な高レベルトレーニングAPI(eager execution前提)。

TF1時代は以下のような違いがありました:

  • tf.Session() や tf.placeholder を明示的に使う
  • tf.estimator や低レベルAPI(tf.train など)が主流
  • Kerasは tf.keras として一部使えましたが、データ入力はNumPyを直接渡すことが多かった
  • model.fit() に tf.data.Dataset を渡すスタイルはTF2で本格的に整備・推奨された

TF1スタイル

# TF1で動くように修正した例
import tensorflow as tf
import numpy as np

(train_images, train_labels), _ = tf.keras.datasets.fashion_mnist.load_data()

train_images = train_images / 255.0
train_labels = train_labels.astype(np.int32)

# tf.data を使わず、NumPy配列を直接渡す(TF1で最も安定)
model = tf.keras.Sequential([
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10)
])

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=2, batch_size=32)

どちらが高速?

  1. 入力パイプラインの違いによる速度比較
  • TF1スタイル(NumPy配列を直接 model.fit(x, y) に渡す)
    • メリット:コードがシンプル。データがメモリにすべて載る小規模データセットではオーバーヘッドが少ない。
    • デメリット:毎エポックでTensorFlowが内部的にデータをコピーしたり、GPUへの転送が非効率になりやすい。特にGPU使用時や複数エポックで入力がボトルネックになる。
  • TF2スタイル(tf.data.Dataset + shuffle + batch)
    • メリット:TensorFlow公式が強く推奨する理由があり、入力パイプラインが大幅に最適化される。
      • 非同期データ読み込み(prefetch)
      • 効率的なバッチ処理とシャッフル
      • GPU/CPU間のデータ転送が最適化されやすい
    • 特にGPUを使っている場合や、エポック数が多いときに入力待ち時間が減り、GPU利用率が向上して全体のトレーニングが速くなる。
  1. 実行モードの影響(TF2特有)TF2のmodel.fit()はデフォルトでeager executionですが、内部的にグラフ最適化(tf.function相当)も一部使っています。
  • 純粋なTF1(静的グラフ + Session)は、シンプルなモデルではまだ高速な場合があります(特に古いベンチマークではTF1が勝つ例あり)。
  • しかし、現代のTF2.10以降では、tf.data + tf.functionの組み合わせがTF1時代の静的グラフを上回るか同等以上の性能を出せるよう最適化が進んでいます。