TensorFlow 2.0 での tf.function と AutoGraph

View on TensorFlow.org Run in Google Colab View source on GitHub

TensorFlow 2.0 では Eager Execution の使いやすさとTensorFlow 1.0 のパワーとを同時に提供します。この統合の中核となるのは tf.function です。これは Python の構文のサブセットを移植可能でハイパフォーマンスな TensorFlow のグラフに変換します。

tf.functionの魅力的な特徴に AutoGraph があります。これはグラフを Python の構文そのものを用いて記述できるようにします。 AutoGraph で利用可能な Python の機能の一覧は、AutoGraph Capabilities and Limitations (Autograph の性能と制限事項) で確認できます。また、tf.functionの詳細については RFC TF 2.0: Functions, not Sessions を参照してください。AutoGraph の詳細については tf.autograph を参照してください。

このチュートリアルでは tf.function と AutoGraph の基本的な特徴についてひととおり確認します。

セットアップ

TensorFlow 2.0 をインポートして、TF 2.0 モードを有効にしてください。

from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
try:
  !pip install -q tf-nightly-2.0-preview
except Exception:
  pass
import tensorflow as tf
ERROR: tensorflow-gpu 2.0.0b1 has requirement tb-nightly<1.14.0a20190604,>=1.14.0a20190603, but you'll have tb-nightly 1.15.0a20190822 which is incompatible.

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=23, shape=(3, 3), dtype=float32, numpy=
array([[0.7202896 , 0.35719517, 0.6748856 ],
       [0.5720995 , 0.4505739 , 1.0261298 ],
       [0.65015596, 0.31933784, 0.6275726 ]], dtype=float32)>

アノテーションの結果を調べてみると、 TensorFlow ランタイムとのやり取りのすべてを処理する特別な呼び出し可能オブジェクトを確認できます。

simple_nn_layer
<tensorflow.python.eager.def_function.Function at 0x7f53824bc9e8>

記述したコードで複数の関数を利用していたとしても、すべての関数にアノテーションを付ける必要はありません。アノテーションをつけた関数から呼び出されるすべての関数は、グラフモードで実行されます。

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=35, shape=(3,), dtype=int32, numpy=array([3, 5, 7], dtype=int32)>

グラフが大量の軽量な演算から構成される場合、関数は Eager Execution で実行するコードよりも高速になる場合があります。しかし、 graph が少量の (畳み込み演算のような) 計算に時間のかかる演算からなる場合、高速化はそれほど見込めないでしょう。

import timeit
conv_layer = tf.keras.layers.Conv2D(100, 3)

@tf.function
def conv_fn(image):
  return conv_layer(image)

image = tf.zeros([1, 200, 200, 100])
# warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Note how there's not much difference in performance for convolutions")
Eager conv: 0.3581459650000056
Function conv: 0.2988417180000056
Note how there's not much difference in performance for convolutions
lstm_cell = tf.keras.layers.LSTMCell(10)

@tf.function
def lstm_fn(input, state):
  return lstm_cell(input, state)

input = tf.zeros([10, 10])
state = [tf.zeros([10, 10])] * 2
# warm up
lstm_cell(input, state); lstm_fn(input, state)
print("eager lstm:", timeit.timeit(lambda: lstm_cell(input, state), number=10))
print("function lstm:", timeit.timeit(lambda: lstm_fn(input, state), number=10))
eager lstm: 0.007858811000005517
function lstm: 0.006511998000007679

Python の制御フローの利用

tf.functionの内部でデータに依存した制御フローを用いる場合、Pythonの制御フロー構文を用いることができます。AutoGraph はそれらの構文を TensorFlow の Ops に書き換えます。たとえば、 Tensor に依存する if 文は、tf.cond() に変換されます。

次の例では xTensor です。ですが、 if 文は期待するどおりに動作しています。

@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

AutoGraph は while, for, if, break, continue, return といった典型的なPythonの構文をサポートしています。また、これらを入れ子にして利用する場合もサポートしています。つまり、Tensorを返す式をwhile 文や if 文の条件式として用いることが可能です。また、for 文で Tensor の要素に渡って反復することも可能です。

@tf.function
def sum_even(items):
  s = 0
  for c in items:
    if c % 2 > 0:
      print(c)
      continue
    s += c
  return s


sum_even(tf.constant([10, 12, 15, 20]))
Tensor("TensorArrayV2Read/TensorListGetItem:0", shape=(), dtype=int32)

<tf.Tensor: id=602, shape=(), dtype=int32, numpy=42>

より高度な使い方をするユーザーのために、AutoGraph は低レベルAPIも提供しています。次の例では AutoGraph が生成したコードを確認できます。

print(tf.autograph.to_code(sum_even.python_function))
def tf__sum_even(items):
  do_return = False
  retval_ = ag__.UndefinedReturnValue()
  with ag__.FunctionScope('sum_even', 'sum_even_scope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as sum_even_scope:
    s = 0

    def get_state_2():
      return ()

    def set_state_2(_):
      pass

    def loop_body(iterates, s):
      c = iterates
      continue_ = False

      def get_state():
        return ()

      def set_state(_):
        pass

      def if_true():
        print(c)
        continue_ = True
        return continue_

      def if_false():
        return continue_
      cond = c % 2 > 0
      continue_ = ag__.if_stmt(cond, if_true, if_false, get_state, set_state, ('continue_',), ())

      def get_state_1():
        return ()

      def set_state_1(_):
        pass

      def if_true_1():
        s_1, = s,
        s_1 += c
        return s_1

      def if_false_1():
        return s
      cond_1 = ag__.not_(continue_)
      s = ag__.if_stmt(cond_1, if_true_1, if_false_1, get_state_1, set_state_1, ('s',), ())
      return s,
    s, = ag__.for_stmt(items, None, loop_body, get_state_2, set_state_2, (s,), ('s',), ())
    do_return = True
    retval_ = sum_even_scope.mark_return_value(s)
  do_return,
  return ag__.retval(retval_)

次はより複雑な制御フローの例です。

@tf.function
def fizzbuzz(n):
  msg = tf.constant('')
  for i in tf.range(n):
    if i % 3 == 0:
      tf.print('Fizz')
    elif i % 5 == 0:
      tf.print('Buzz')
    else:
      tf.print(i)

fizzbuzz(tf.constant(15))
Fizz
1
2
Fizz
4
Buzz
Fizz
7
8
Fizz
Buzz
11
Fizz
13
14

Keras での AutoGraph の利用

tf.function はオブジェクトのメソッドに対しても利用できます。たとえば、カスタムしたKeras モデルにデコレーターを適用できます、典型的には call 関数にアノテーションを付けることで実現できるでしょう。より詳細が必要な場合、tf.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=711, shape=(2,), dtype=int32, numpy=array([-1, -2], dtype=int32)>

副作用

Eager モードのように、通常の場合 tf.function の中で、tf.assigntf.print といった副作用のある命令を実行できます。また、実行時の順序を保つために、処理順について必要な依存関係を書き加えます。

v = tf.Variable(5)

@tf.function
def find_next_odd():
  v.assign(v + 1)
  if 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 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.7020334 ; accuracy 0.374
Step 20 : loss 1.20663249 ; accuracy 0.521
Step 30 : loss 0.661244333 ; accuracy 0.607333362
Step 40 : loss 0.63975656 ; accuracy 0.668
Step 50 : loss 0.458514124 ; accuracy 0.7058
Step 60 : loss 0.441733211 ; accuracy 0.733333349
Step 70 : loss 0.454216689 ; accuracy 0.752428591
Step 80 : loss 0.411001056 ; accuracy 0.770375
Step 90 : loss 0.300666898 ; accuracy 0.784777761
Step 100 : loss 0.358992308 ; accuracy 0.7976
Step 110 : loss 0.317251444 ; accuracy 0.805909097
Step 120 : loss 0.359570056 ; accuracy 0.813083351
Step 130 : loss 0.249649644 ; accuracy 0.82076925
Step 140 : loss 0.287468493 ; accuracy 0.827642858
Step 150 : loss 0.327236414 ; accuracy 0.8326
Step 160 : loss 0.410696745 ; accuracy 0.837812483
Step 170 : loss 0.303566545 ; accuracy 0.841470599
Step 180 : loss 0.260952801 ; accuracy 0.84577775
Step 190 : loss 0.285236478 ; accuracy 0.850157917
Step 200 : loss 0.299665421 ; accuracy 0.85275
Final step tf.Tensor(200, shape=(), dtype=int32) : loss tf.Tensor(0.29966542, shape=(), dtype=float32) ; accuracy tf.Tensor(0.85275, 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=1661, 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=1670, shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1,  0,  1,  4,  9, 16], dtype=int32)>