TensorFlow 2.0 Alpha : ガイド : TensorFlow 2.0 の tf.function と AutoGraph (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/13/2019
* 本ページは、TensorFlow の本家サイトの TF 2.0 Alpha の以下のページを翻訳した上で適宜、補足説明したものです：
* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
ガイド : TensorFlow 2.0 の tf.function と AutoGraph
TF 2.0 は eager execution の容易さと TF 1.0 のパワーを一つにまとめています。この融合の中心には tf.function があります、これは Python シンタクスのサブセットを可搬な、高パフォーマンスの TensorFlow グラフに変換することを可能にします。
tf.function のクールな新しい特徴は AutoGraph です、これは貴方に自然な Python シンタクスを使用してグラフコードを書かせます。AutoGraph で使用できる Python 特徴のリストについては、AutoGraph Capabilities and Limitations を見てください。tf.function についてのより多くの詳細は、RFC TF 2.0: Functions, not Sessions を見てください。このチュートリアルは tf.function と AutoGraph の基本的な特徴を貴方にウォークスルーしてもらいます。
Setup
TensorFlow 2.0 Preview Nightly をインポートして TF 2.0 モードを有効にします :
from __future__ import absolute_import, division, print_function import numpy as np
!pip install -q tensorflow==2.0.0-alpha0 import tensorflow as tf
幾つかの特別な TF 2.0 アップグレードを有効にするために一時パッチをインストールします。このピースは間もなく除去されるでしょう。
from tensorflow.python.ops import control_flow_util control_flow_util.ENABLE_CONTROL_FLOW_V2 = True
tf.function デコレータ
関数を tf.function でアノテートするとき、依然として任意の他の関数のようにそれを呼び出すことができます。しかしそれはグラフ内にコンパイルされ、これは高速な実行の恩恵を得ることを意味し、GPU か TPU で実行されるか、SavedModel にエクスポートされます。
@tf.function def simple_nn_layer(x, y): return tf.nn.relu(tf.matmul(x, y)) x = tf.random.uniform((3, 3)) y = tf.random.uniform((3, 3)) simple_nn_layer(x, y)
<tf.Tensor: id=25, shape=(3, 3), dtype=float32, numpy=
array([[0.29857793, 0.3464706 , 0.56053376],
[0.3308559 , 0.30354396, 0.5458242 ],
[0.7290289 , 0.43530926, 0.9764521 ]], dtype=float32)>
アノテーションの結果を調べればそれが TensorFlow ランタイムとの総ての相互作用を扱う特別な callable であることが見て取れるでしょう。
simple_nn_layer
<tensorflow.python.eager.def_function.Function at 0x7fb8ae4629e8>
貴方のコードが複数の関数を使用する場合、それら総てをアノテートする必要はありません – アノテートされた関数から呼び出された任意の関数もまたグラフモードで動作します。
def linear_layer(x): return 2 * x + 1 @tf.function def deep_net(x): return tf.nn.relu(linear_layer(x)) deep_net(tf.constant((1, 2, 3)))
<tf.Tensor: id=39, shape=(3,), dtype=int32, numpy=array([3, 5, 7], dtype=int32)>
Python 制御フローを使用する
tf.function 内でデータ依存制御フローを使用しているとき、Python 制御フロー・ステートメントを使用することができてそして AutoGraph はそれらを適切な TensorFlow ops に変換します。例えば、ステートメントはそれらが Tensor に依存する場合 tf.cond() に変換されるでしょう。下の例では、x は Tensor ですがステートメントは期待通りに動作します :
@tf.function
def square_if_positive(x):
if x > 0:
x = x * x
else:
x = 0
return x
print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))
print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))
square_if_positive(2) = 4 square_if_positive(-2) = 0
Note: 上の例はスカラー値が関係するときに単一の条件をどのように遂行するかを示しています。典型的な ML コードはバッチを伴います ; その場合には可能であればより高速でベクトル化された tf.where を使用することを考慮するべきです。
AutoGraph は while, for, if, break, continue と return のような一般的な Python ステートメントをネスティングのためのサポートとともに、サポートします。それは while と if ステートメントの条件で Tensor 式を使用したり、for ループで Tensor に渡り iterate できることを意味します。
@tf.function
def sum_even(items):
s = 0
for c in items:
if c % 2 > 0:
continue
s += c
return s
sum_even(tf.constant([10, 12, 15, 20]))
<tf.Tensor: id=149, shape=(), dtype=int32, numpy=42>
AutoGraph はまた上級ユーザのために低位 API も提供します。例えば生成されたコードを見るためにそれを使用できます。
print(tf.autograph.to_code(sum_even.python_function, experimental_optional_features=None))
from __future__ import print_function
def tf__sum_even(items):
do_return = False
retval_ = None
s = 0
def loop_body(loop_vars, s_2):
c = loop_vars
continue_ = False
cond = c % 2 > 0
def if_true():
continue_ = True
return continue_
def if_false():
return continue_
continue_ = ag__.if_stmt(cond, if_true, if_false)
cond_1 = ag__.not_(continue_)
def if_true_1():
s_1, = s_2,
s_1 += c
return s_1
def if_false_1():
return s_2
s_2 = ag__.if_stmt(cond_1, if_true_1, if_false_1)
return s_2,
s, = ag__.for_stmt(items, None, loop_body, (s,))
do_return = True
retval_ = s
return retval_
tf__sum_even.autograph_info__ = {}
ここにより複雑な制御フローのサンプルがあります :
@tf.function
def fizzbuzz(n):
msg = tf.constant('')
for i in tf.range(n):
if tf.equal(i % 3, 0):
msg += 'Fizz'
elif tf.equal(i % 5, 0):
msg += 'Buzz'
else:
msg += tf.as_string(i)
msg += '\n'
return msg
print(fizzbuzz(tf.constant(15)).numpy().decode())
Fizz 1 2 Fizz 4 Buzz Fizz 7 8 Fizz Buzz 11 Fizz 13 14
Keras と AutoGraph
tf.function はオブジェクトのメソッドでもまた使用できます。例えば、典型的にはモデルの call 関数をアノテートすることで貴方のカスタム Keras モデルをデコレートできます。
class CustomModel(tf.keras.models.Model):
@tf.function
def call(self, input_data):
if tf.reduce_mean(input_data) > 0:
return input_data
else:
return input_data // 2
model = CustomModel()
model(tf.constant([-2, -4]))
<tf.Tensor: id=281, shape=(2,), dtype=int32, numpy=array([-1, -2], dtype=int32)>
副作用
ちょうど eager モードでのように、通常は tf.function 内で tf.assign や tf.print のような、副作用を持つ演算を使用することができて、それはそれらが順番に実行されることを確実にするために必要な制御依存性を挿入します。
v = tf.Variable(5)
@tf.function
def find_next_odd():
v.assign(v + 1)
if tf.equal(v % 2, 0):
v.assign(v + 1)
find_next_odd()
v
<tf.Variable 'Variable:0' shape=() dtype=int32, numpy=7>
サンプル: 単純なモデルを訓練する
AutoGraph はまたより多くの計算を TensorFlow 内に移すことを可能にします。例えば、訓練ループはちょうど制御フローですので、それは実際には TensorFlow に持ち込むことができます。
データをダウンロードする
def prepare_mnist_features_and_labels(x, y): x = tf.cast(x, tf.float32) / 255.0 y = tf.cast(y, tf.int64) return x, y def mnist_dataset(): (x, y), _ = tf.keras.datasets.mnist.load_data() ds = tf.data.Dataset.from_tensor_slices((x, y)) ds = ds.map(prepare_mnist_features_and_labels) ds = ds.take(20000).shuffle(20000).batch(100) return ds train_dataset = mnist_dataset()
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz 11493376/11490434 [==============================] - 0s 0us/step
モデルを定義する
model = tf.keras.Sequential((
tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),
tf.keras.layers.Dense(100, activation='relu'),
tf.keras.layers.Dense(100, activation='relu'),
tf.keras.layers.Dense(10)))
model.build()
optimizer = tf.keras.optimizers.Adam()
訓練ループを定義する
compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
def train_one_step(model, optimizer, x, y):
with tf.GradientTape() as tape:
logits = model(x)
loss = compute_loss(y, logits)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
compute_accuracy(y, logits)
return loss
@tf.function
def train(model, optimizer):
train_ds = mnist_dataset()
step = 0
loss = 0.0
accuracy = 0.0
for x, y in train_ds:
step += 1
loss = train_one_step(model, optimizer, x, y)
if tf.equal(step % 10, 0):
tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
return step, loss, accuracy
step, loss, accuracy = train(model, optimizer)
print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
Step 10 : loss 1.72189403 ; accuracy 0.378 Step 20 : loss 1.17039371 ; accuracy 0.5105 Step 30 : loss 0.858835042 ; accuracy 0.600666642 Step 40 : loss 0.501416266 ; accuracy 0.659 Step 50 : loss 0.501890719 ; accuracy 0.6984 Step 60 : loss 0.402025729 ; accuracy 0.723 Step 70 : loss 0.381988674 ; accuracy 0.747571409 Step 80 : loss 0.287695855 ; accuracy 0.76675 Step 90 : loss 0.388588548 ; accuracy 0.780111134 Step 100 : loss 0.280551046 ; accuracy 0.7926 Step 110 : loss 0.251829863 ; accuracy 0.803181827 Step 120 : loss 0.271081507 ; accuracy 0.812166691 Step 130 : loss 0.264373362 ; accuracy 0.819615364 Step 140 : loss 0.377884507 ; accuracy 0.826071441 Step 150 : loss 0.326621652 ; accuracy 0.8312 Step 160 : loss 0.470905215 ; accuracy 0.836375 Step 170 : loss 0.410217643 ; accuracy 0.840647042 Step 180 : loss 0.393875092 ; accuracy 0.844055533 Step 190 : loss 0.283068538 ; accuracy 0.848157883 Step 200 : loss 0.20351316 ; accuracy 0.8513 Final step tf.Tensor(200, shape=(), dtype=int32) : loss tf.Tensor(0.20351316, shape=(), dtype=float32) ; accuracy tf.Tensor(0.8513, shape=(), dtype=float32)
バッチ処理についてのノート
実際のアプリケーションではバッチ処理はパフォーマンスのために重要です。AutoGraph に変換する最善のコードはそこでは制御フローがバッチレベルで決定されるようなコードです。個々のサンプルレベルで決定を行なう場合には、パフォーマンスを維持するために batch API を利用してみてください。
例えば、Python で次のようなコードを持つ場合 :
def square_if_positive(x): return [i ** 2 if i > 0 else i for i in x] square_if_positive(range(-5, 5))
[-5, -4, -3, -2, -1, 0, 1, 4, 9, 16]
TensorFlow で次のようにそれを書く誘惑にかられるかもしれません (そしてこれは動作します！) :
@tf.function
def square_if_positive_naive(x):
result = tf.TensorArray(tf.int32, size=x.shape[0])
for i in tf.range(x.shape[0]):
if x[i] > 0:
result = result.write(i, x[i] ** 2)
else:
result = result.write(i, x[i])
return result.stack()
square_if_positive_naive(tf.range(-5, 5))
<tf.Tensor: id=1544, shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1, 0, 1, 4, 9, 16], dtype=int32)>
しかしこの場合、次を書くことができることが判明します :
def square_if_positive_vectorized(x): return tf.where(x > 0, x ** 2, x) square_if_positive_vectorized(tf.range(-5, 5))
<tf.Tensor: id=1554, shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1, 0, 1, 4, 9, 16], dtype=int32)>
以上