TensorFlow 2.0 : 上級 Tutorials : カスタマイズ :- tf.function でパフォーマンスの改善 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 10/29/2019

* 本ページは、TensorFlow org サイトの TF 2.0 – Advanced Tutorials – Customization の以下のページを翻訳した上で適宜、補足説明したものです：

• Better performance with tf.function

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

カスタマイズ :- tf.function でパフォーマンスの改善

TensorFlow 2.0 では eager execution はデフォルトで有効になっています。ユーザインターフェイスは直感的で柔軟ですが (一度限りの (= one-off) 演算の実行は遥かに容易で高速です)、これはパフォーマンスと配備性 (= deployability) の犠牲をもたらす可能性があります。

最高点のパフォーマンスを得て貴方のモデルをどこでも配備可能にするために、貴方のプログラムからグラフを作成するために tf.function を使用します。AutoGraph のおかげで、Python コードの驚くべき総量が tf.function で単に動作しますが、依然として注意すべき落とし穴があります。

主な重要な点と推奨は :

• オブジェクト mutation やリスト append のような Python 副作用に依拠しないでください。
• tf.function は NumPy ops や Python プリミティブよりも TensorFlow ops で最善に動作します。
• 疑わしいときには、”for x in y” イディオムを使用してください。

from __future__ import absolute_import, division, print_function, unicode_literals

import tensorflow as tf

import contextlib

# Some helper code to demonstrate the kinds of errors you might encounter.
@contextlib.contextmanager
def assert_raises(error_class):
  try:
    yield
  except error_class as e:
    print('Caught expected exception \n  {}: {}'.format(error_class, e))
  except Exception as e:
    print('Got unexpected exception \n  {}: {}'.format(type(e), e))
  else:
    raise Exception('Expected {} to be raised but no error was raised!'.format(
        error_class))

貴方が定義する tf.function はちょうどコア TensorFlow 演算のようなものです : それを eagerly (逐次実行的) に実行できます; それをグラフで使用できます; それは勾配を持ちます; 等々。

# A function is like an op

@tf.function
def add(a, b):
  return a + b

add(tf.ones([2, 2]), tf.ones([2, 2]))  #  [[2., 2.], [2., 2.]]

<tf.Tensor: id=14, shape=(2, 2), dtype=float32, numpy=
array([[2., 2.],
       [2., 2.]], dtype=float32)>

# Functions have gradients

@tf.function
def add(a, b):
  return a + b

v = tf.Variable(1.0)
with tf.GradientTape() as tape:
  result = add(v, 1.0)
tape.gradient(result, v)

<tf.Tensor: id=38, shape=(), dtype=float32, numpy=1.0>

# You can use functions inside functions

@tf.function
def dense_layer(x, w, b):
  return add(tf.matmul(x, w), b)

dense_layer(tf.ones([3, 2]), tf.ones([2, 2]), tf.ones([2]))

<tf.Tensor: id=64, shape=(3, 2), dtype=float32, numpy=
array([[3., 3.],
       [3., 3.],
       [3., 3.]], dtype=float32)>

 

Tracing とポリモーフィズム

Python の動的型付けは関数を様々な引数型で呼び出すことができて、そして Python は各シナリオで異なる何かを行なうことを意味します。

その一方で、TensorFlow グラフは静的な dtype と shape 次元を要求します。tf.function は正しいグラフを生成するために必要なとき関数を辿り直す (= retrace) ことによりこの隔たりの橋渡しをします。tf.function の利用方法の微妙な点の殆どはこの再トレースの動作から生じています。

何が起きるかを見るために異なる型の引数で関数を呼び出すことができます。

# Functions are polymorphic

@tf.function
def double(a):
  print("Tracing with", a)
  return a + a

print(double(tf.constant(1)))
print()
print(double(tf.constant(1.1)))
print()
print(double(tf.constant("a")))
print()

Tracing with Tensor("a:0", shape=(), dtype=int32)
tf.Tensor(2, shape=(), dtype=int32)

Tracing with Tensor("a:0", shape=(), dtype=float32)
tf.Tensor(2.2, shape=(), dtype=float32)

Tracing with Tensor("a:0", shape=(), dtype=string)
tf.Tensor(b'aa', shape=(), dtype=string)

トレース動作を制御するには、次のテクニックを使用します :

• 新しい tf.function を作成します。別々の tf.function オブジェクトはトレースを共有しないことが保証されます。
• 特定のトレース を得るために get_concrete_function メソッドを使用します。
• グラフ呼び出し毎に一度だけトレースするために tf.function を呼び出すとき input_signature を指定します。

print("Obtaining concrete trace")
double_strings = double.get_concrete_function(tf.TensorSpec(shape=None, dtype=tf.string))
print("Executing traced function")
print(double_strings(tf.constant("a")))
print(double_strings(a=tf.constant("b")))
print("Using a concrete trace with incompatible types will throw an error")
with assert_raises(tf.errors.InvalidArgumentError):
  double_strings(tf.constant(1))

Obtaining concrete trace
Tracing with Tensor("a:0", dtype=string)
Executing traced function
tf.Tensor(b'aa', shape=(), dtype=string)
tf.Tensor(b'bb', shape=(), dtype=string)
Using a concrete trace with incompatible types will throw an error
Caught expected exception 
  <class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>: cannot compute __inference_double_91 as input #0(zero-based) was expected to be a string tensor but is a int32 tensor [Op:__inference_double_91]

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(x % 2 == 0, x // 2, 3 * x + 1)

print(next_collatz(tf.constant([1, 2])))
# We specified a 1-D tensor in the input signature, so this should fail.
with assert_raises(ValueError):
  next_collatz(tf.constant([[1, 2], [3, 4]]))

Tracing with Tensor("x:0", shape=(None,), dtype=int32)
tf.Tensor([4 1], shape=(2,), dtype=int32)
Caught expected exception 
  <class 'ValueError'>: Python inputs incompatible with input_signature:
  inputs: (
    tf.Tensor(
[[1 2]
 [3 4]], shape=(2, 2), dtype=int32))
  input_signature: (
    TensorSpec(shape=(None,), dtype=tf.int32, name=None))

 

いつ再トレースするか？

多態的な (= polymorphic) tf.function はトレースにより生成された具体的な関数のキャッシュを保持します。キャッシュ・キーは効果的に関数 argsと kwargs から生成されたキーのタプルです。tf.Tensor 引数のために生成されたキーはその shape と型です。Python プリミティブのために生成されたキーはその値です。総ての他の Python 型については、キーはオブジェクト id() に基づきその結果メソッドはクラスの各インスタンスのために独立的にトレースされます。将来的には、TensorFlow は Python オブジェクトのためにより洗練されたキャッシングを追加するかもしれません、それは tensor に安全に変換できます。

 

Python あるいは Tensor args?

しばしば、Python 引数はハイパーパラメータとグラフ構築を制御するために使用されます – 例えば、num_layers=10 や training=True や nonlinearity=’relu’ です。そのため Python 引数が変わる場合、グラフを再トレースしなければならないことは意味があります。

けれども、Python 引数がグラフ構築を制御するために使用されていないこともあり得ます。これらのケースでは、Python 値の変更は必要のない再トレースを引き起こす可能性があります。例えばこの訓練ループを取れば、AutoGraph はそれを動的に展開するでしょう。複数のトレースにもかかわらず、生成されたグラフは実際には同一です、従ってこれは少し非効率的です。

def train_one_step():
  pass

@tf.function
def train(num_steps):
  print("Tracing with num_steps = {}".format(num_steps))
  for _ in tf.range(num_steps):
    train_one_step()

train(num_steps=10)
train(num_steps=20)

Tracing with num_steps = 10
Tracing with num_steps = 20

ここでの最も単純な回避方法は、それら (引数) が生成されたグラフの shape に影響しないのであれば、引数を Tensor にキャストすることです。

train(num_steps=tf.constant(10))
train(num_steps=tf.constant(20))

Tracing with num_steps = Tensor("num_steps:0", shape=(), dtype=int32)

 

tf.function の副作用

一般に、(プリントやオブジェクトの mutating のような) Python 副作用はトレースの間だけ発生します。それでは tf.function からどのように副作用を確実に引き起こせるのでしょう？

一般的な経験則はトレースをデバッグするために Python 副作用を使用するだけです。そうでないなら、 tf.Variable.assign, tf.print と tf.summary のような TensorFlow ops は各呼び出しに伴う TensorFlow ランタイムにより貴方のコードがトレースされて実行されることを確かなものにする最善の方法です。一般に functional スタイルの使用は最善の結果を生じるでしょう。

@tf.function
def f(x):
  print("Traced with", x)
  tf.print("Executed with", x)

f(1)
f(1)
f(2)

Traced with 1
Executed with 1
Executed with 1
Traced with 2
Executed with 2

tf.function の各起動中に Python コードを実行することを望むのであれば、tf.py_function が脱出ハッチ (= exit hatch) です。tf.py_function の欠点はそれは可搬ではなく特に非効率的で、分散セットアップ (マルチ GPU, TPU) では上手く動作しません。また、tf.py_function はグラフに配線 (= wire) されなければならず、それは総ての入力/出力を tensor にキャストします。

external_list = []

def side_effect(x):
  print('Python side effect')
  external_list.append(x)

@tf.function
def f(x):
  tf.py_function(side_effect, inp=[x], Tout=[])

f(1)
f(1)
f(1)
assert len(external_list) == 3
# .numpy() call required because py_function casts 1 to tf.constant(1)
assert external_list[0].numpy() == 1

Python side effect
Python side effect
Python side effect

 

Python 状態に注意する

generator と iterator のような、多くの Python 特徴は状態を追跡するために Python ランタイムに依拠します。一般に、これらの構成物が Eager モードで期待どおりに動作する一方で、多くの予期せぬことがトレース動作により tf.function 内で起きるかもしれません。

一つのサンプルを与えるために、iterator 状態を前に進めることは Python 副作用で従ってトレースの間だけに発生します。

external_var = tf.Variable(0)
@tf.function
def buggy_consume_next(iterator):
  external_var.assign_add(next(iterator))
  tf.print("Value of external_var:", external_var)

iterator = iter([0, 1, 2, 3])
buggy_consume_next(iterator)
# This reuses the first value from the iterator, rather than consuming the next value.
buggy_consume_next(iterator)
buggy_consume_next(iterator)

Value of external_var: 0
Value of external_var: 0
Value of external_var: 0

iterator が完全に tf.function 内で生成されて消費される場合には、それは正しく動作します。けれども、iterator 全体がトレースされると多分、これは巨大グラフ (= giant graph) に繋がるかもしれません。これは貴方が望むものかもしれません。しかし貴方が Python リストで表わされる巨大な in-memory データセット上で訓練している場合、これは非常に巨大なグラフを生成するかもしれず、そして tf.function はスピードアップを生み出さないでしょう。

Python データに渡り iterate することを望む場合、最も安全な方法はそれを tf.data.Dataset でラップして “for x in y” イディオムを使用することです。AutoGraph は y が tensor か tf.data.Dataset のときループのための安全に変換するための特別なサポートを持ちます。

def measure_graph_size(f, *args):
  g = f.get_concrete_function(*args).graph
  print("{}({}) contains {} nodes in its graph".format(
      f.__name__, ', '.join(map(str, args)), len(g.as_graph_def().node)))

@tf.function
def train(dataset):
  loss = tf.constant(0)
  for x, y in dataset:
    loss += tf.abs(y - x) # Some dummy computation.
  return loss

small_data = [(1, 1)] * 2
big_data = [(1, 1)] * 10
measure_graph_size(train, small_data)
measure_graph_size(train, big_data)

measure_graph_size(train, tf.data.Dataset.from_generator(
    lambda: small_data, (tf.int32, tf.int32)))
measure_graph_size(train, tf.data.Dataset.from_generator(
    lambda: big_data, (tf.int32, tf.int32)))

train([(1, 1), (1, 1)]) contains 8 nodes in its graph
train([(1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1), (1, 1)]) contains 32 nodes in its graph
train(, ), types: (tf.int32, tf.int32)>) contains 5 nodes in its graph
train(, ), types: (tf.int32, tf.int32)>) contains 5 nodes in its graph

Python/Numpy データを Dataset にラッピングするとき、tf.data.Dataset.from_generator versus tf.data.Dataset.from_tensors に留意してください。前者はデータを Python 内で保持して tf.py_function を通してそれを取得します、これはパフォーマンスと密接な関係を持つかもしれず、その一方で後者はデータのコピーをグラフの一つの巨大な tf.constant() ノードとして束ねます、これはメモリとの密接な関係を持つかもしれません。

TFRecordDataset/CsvDataset/etc. を通してファイルからデータを読むのがデータを消費する最も効果的な方法です、何故ならばそれで TensorFlow 自身がデータの非同期ロードと先読みを管理できるからです、Python を巻き込まなければならないことなしに。

 

自動制御依存性

functions のプログラミングモデルとしての非常に魅力的な特質は、一般的なデータフロー・グラフに渡り、functions がランタイムに何がコードの意図された動作であったかについてより多くの情報を与えられることです。

例えば、同じ変数への複数の読み書きを持つコードを書くとき、データフロー・グラフは演算の元々意図された順序を自然にエンコードしないかもしれません。tf.function では、元の Python コードのステートメントの実行順序を参照して実行順序の曖昧さを解決します。このようにして、tf.function のステートフルな演算の順序は Eager モードのセマンティクスを複製します。

これは手動の制御依存性を追加する必要がないことを意味します ; tf.function は貴方のコードが正しく実行されるために必要で十分な制御依存性の最小セットを追加するために十分にスマートです。

# Automatic control dependencies

a = tf.Variable(1.0)
b = tf.Variable(2.0)

@tf.function
def f(x, y):
  a.assign(y * b)
  b.assign_add(x * a)
  return a + b

f(1.0, 2.0)  # 10.0

<tf.Tensor: id=418, shape=(), dtype=float32, numpy=10.0>

 

Variables

tf.function で変数作成と利用を非常に容易にするためにコードの意図された実行順序を活用するのと同じアイデアを利用できます。けれども一つの非常に重要な警告があります、それは variable では eager モードと graph モードで異なる動作をするコードを書くことが可能であることです。

特に、これは各呼び出しで新しい Variable を作成するときに発生します。トレース・セマンティクスにより、tf.function は各呼び出しで同じ変数を再利用しますが、eager モードは各呼び出しで新しい変数を作成します。この間違いから守るために、 tf.function は危険な変数作成動作を検出する場合にはエラーを上げます。

@tf.function
def f(x):
  v = tf.Variable(1.0)
  v.assign_add(x)
  return v

with assert_raises(ValueError):
  f(1.0)

WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1781: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
Caught expected exception 
  <class 'ValueError'>: in converted code:

    <ipython-input-17-73e410646579>:3 f  *
        v = tf.Variable(1.0)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/variables.py:260 __call__
        return cls._variable_v2_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/variables.py:254 _variable_v2_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/variables.py:65 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/eager/def_function.py:413 invalid_creator_scope
        "tf.function-decorated function tried to create "

    ValueError: tf.function-decorated function tried to create variables on non-first call.

# Non-ambiguous code is ok though

v = tf.Variable(1.0)

@tf.function
def f(x):
  return v.assign_add(x)

print(f(1.0))  # 2.0
print(f(2.0))  # 4.0

tf.Tensor(2.0, shape=(), dtype=float32)
tf.Tensor(4.0, shape=(), dtype=float32)

# You can also create variables inside a tf.function as long as we can prove
# that those variables are created only the first time the function is executed.

class C: pass
obj = C(); obj.v = None

@tf.function
def g(x):
  if obj.v is None:
    obj.v = tf.Variable(1.0)
  return obj.v.assign_add(x)

print(g(1.0))  # 2.0
print(g(2.0))  # 4.0

tf.Tensor(2.0, shape=(), dtype=float32)
tf.Tensor(4.0, shape=(), dtype=float32)

# Variable initializers can depend on function arguments and on values of other
# variables. We can figure out the right initialization order using the same
# method we use to generate control dependencies.

state = []
@tf.function
def fn(x):
  if not state:
    state.append(tf.Variable(2.0 * x))
    state.append(tf.Variable(state[0] * 3.0))
  return state[0] * x * state[1]

print(fn(tf.constant(1.0)))
print(fn(tf.constant(3.0)))

tf.Tensor(12.0, shape=(), dtype=float32)
tf.Tensor(36.0, shape=(), dtype=float32)

 

AutoGraph を使用する

autograph ライブラリは tf.function と完全に統合され、そしてそれはグラフで動的に実行するための Tensor に依拠する条件節とループを書き換えます。

tf.cond と tf.while_loop は tf.function とともに動作し続けますが、制御フローを持つコードはしばしば命令型スタイルで書かれるとき書いて理解することが容易です。

# Simple loop

@tf.function
def f(x):
  while tf.reduce_sum(x) > 1:
    tf.print(x)
    x = tf.tanh(x)
  return x

f(tf.random.uniform([5]))

[0.991675735 0.770152 0.265438318 0.926479578 0.270098567]
[0.758075953 0.647017837 0.259375095 0.728948355 0.263716549]
[0.639942288 0.569659 0.253710955 0.622421563 0.257768452]
[0.564860284 0.515108824 0.248403832 0.552811861 0.2522071]
[0.511574626 0.473916173 0.24341765 0.50262475 0.246992245]
[0.471171141 0.441358089 0.238721251 0.46417886 0.242089257]
[0.439145088 0.41476953 0.23428756 0.433484 0.237468168]
[0.412935585 0.392514914 0.230092898 0.408228815 0.233102903]
[0.390962422 0.373526275 0.226116508 0.386967748 0.228970662]
[0.372189641 0.357072234 0.222340047 0.368743211 0.225051373]
[0.355905473 0.342632562 0.218747273 0.352891922 0.22132732]
[0.341602385 0.32982561 0.215323746 0.33893773 0.217782795]
[0.328907162 0.318364054 0.212056547 0.326528698 0.214403793]
[0.31753847 0.30802694 0.208934113 0.315398216 0.211177796]
[0.307279497 0.298641056 0.205946043 0.305340052 0.208093569]
[0.297960132 0.290068477 0.203082964 0.296191841 0.205141023]
[0.289444745 0.282197833 0.200336367 0.287823766 0.202311009]
[0.281623662 0.274938 0.197698563 0.280130565 0.199595287]
[0.274407148 0.26821363 0.19516255 0.2730259 0.196986347]
[0.267720908 0.261961818 0.192721918 0.266438186 0.194477364]
[0.261502862 0.256129563 0.190370843 0.260307461 0.192062095]
[0.255700678 0.250671834 0.188103959 0.254583091 0.189734846]
[0.25026986 0.245550096 0.185916349 0.249221966 0.187490389]
[0.245172322 0.24073115 0.183803499 0.244187161 0.185323924]
[0.240375236 0.236186236 0.181761235 0.239446774 0.183231026]
[0.23585014 0.231890276 0.179785714 0.234973133 0.181207627]
[0.231572226 0.22782129 0.177873373 0.230742082 0.179249957]
[0.227519721 0.223959938 0.176020905 0.226732403 0.177354515]
[0.223673478 0.220289111 0.174225256 0.22292541 0.175518081]
[0.220016524 0.216793597 0.172483563 0.219304562 0.173737645]

<tf.Tensor: id=675, shape=(5,), dtype=float32, numpy=
array([0.21653381, 0.21345986, 0.17079319, 0.21585512, 0.17201042],
      dtype=float32)>

# If you're curious you can inspect the code autograph generates.
# It feels like reading assembly language, though.

def f(x):
  while tf.reduce_sum(x) > 1:
    tf.print(x)
    x = tf.tanh(x)
  return x

print(tf.autograph.to_code(f))

def tf__f(x):
  do_return = False
  retval_ = ag__.UndefinedReturnValue()
  with ag__.FunctionScope('f', 'f_scope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as f_scope:

    def get_state():
      return ()

    def set_state(_):
      pass

    def loop_body(x):
      ag__.converted_call(tf.print, f_scope.callopts, (x,), None, f_scope)
      x = ag__.converted_call(tf.tanh, f_scope.callopts, (x,), None, f_scope)
      return x,

    def loop_test(x):
      return ag__.converted_call(tf.reduce_sum, f_scope.callopts, (x,), None, f_scope) > 1
    x, = ag__.while_stmt(loop_test, loop_body, get_state, set_state, (x,), ('x',), ())
    do_return = True
    retval_ = f_scope.mark_return_value(x)
  do_return,
  return ag__.retval(retval_)

 

AutoGraph: 条件節

AutoGraph は if ステートメントを等値な tf.cond 呼び出しに変換します。

この置き換えは条件が Tensor である場合に行われます。そうでないなら、条件節はトレースの間に実行されます。

def test_tf_cond(f, *args):
  g = f.get_concrete_function(*args).graph
  if any(node.name == 'cond' for node in g.as_graph_def().node):
    print("{}({}) uses tf.cond.".format(
        f.__name__, ', '.join(map(str, args))))
  else:
    print("{}({}) executes normally.".format(
        f.__name__, ', '.join(map(str, args))))

@tf.function
def hyperparam_cond(x, training=True):
  if training:
    x = tf.nn.dropout(x, rate=0.5)
  return x

@tf.function
def maybe_tensor_cond(x):
  if x < 0:
    x = -x
  return x

test_tf_cond(hyperparam_cond, tf.ones([1], dtype=tf.float32))
test_tf_cond(maybe_tensor_cond, tf.constant(-1))
test_tf_cond(maybe_tensor_cond, -1)

hyperparam_cond(tf.Tensor([1.], shape=(1,), dtype=float32)) executes normally.
maybe_tensor_cond(tf.Tensor(-1, shape=(), dtype=int32)) uses tf.cond.
maybe_tensor_cond(-1) executes normally.

tf.cond は多くの微妙な点を持ちます。- それは条件節の両側をトレースし、それから条件に依拠して、ランタイムに適切な分岐を選択することにより動作します。両側のトレースは Python コードの予期せぬ実行という結果になるかもしれません - 一つの分岐が downstream で使用される tensor を作成する場合、他の分岐もまた tensor を作成しなければならないことを必要とします。

@tf.function
def f():
  x = tf.constant(0)
  if tf.constant(True):
    x = x + 1
    print("Tracing `then` branch")
  else:
    x = x - 1
    print("Tracing `else` branch")
  return x

f()

Tracing `then` branch
Tracing `else` branch

<tf.Tensor: id=747, shape=(), dtype=int32, numpy=1>

@tf.function
def f():
  if tf.constant(True):
    x = tf.ones([3, 3])
  return x

# Throws an error because both branches need to define `x`.
with assert_raises(ValueError):
  f()

Caught expected exception 
  <class 'ValueError'>: in converted code:

    <ipython-input-26-3ae6a92ff50d>:3 f  *
        if tf.constant(True):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:893 if_stmt
        basic_symbol_names, composite_symbol_names)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:931 tf_if_stmt
        error_checking_orelse)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/util/deprecation.py:507 new_func
        return func(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/control_flow_ops.py:1174 cond
        return cond_v2.cond_v2(pred, true_fn, false_fn, name)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/cond_v2.py:91 cond_v2
        op_return_value=pred)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/func_graph.py:915 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:924 error_checking_orelse
        result[orelse_branch] = orelse()
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:962 wrapper
        new_vars = func()
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:988 wrapper
        tuple(s.symbol_name for s in undefined)))

    ValueError: The following symbols must also be initialized in the else branch: ('x',). Alternatively, you may initialize them before the if statement.

 

AutoGraph とループ

AutoGraph はループを変換するための 2, 3 の単純なルールを持ちます。

• for: iterable が tensor である場合に変換します。
• while: while 条件が tensor に依拠する場合に変換されます。

ループが変換される場合、それは tf.while_loop で動的に変換され、あるいは "for x in tf.data.Dataset" の特別なケースでは、tf.data.Dataset.reduce に変換されます。

ループが変換されない場合には、それは静的に展開されます。

def test_dynamically_unrolled(f, *args):
  g = f.get_concrete_function(*args).graph
  if any(node.name == 'while' for node in g.as_graph_def().node):
    print("{}({}) uses tf.while_loop.".format(
        f.__name__, ', '.join(map(str, args))))
  elif any(node.name == 'ReduceDataset' for node in g.as_graph_def().node):
    print("{}({}) uses tf.data.Dataset.reduce.".format(
        f.__name__, ', '.join(map(str, args))))
  else:
    print("{}({}) gets unrolled.".format(
        f.__name__, ', '.join(map(str, args))))

@tf.function
def for_in_range():
  x = 0
  for i in range(5):
    x += i
  return x

test_dynamically_unrolled(for_in_range)

for_in_range() gets unrolled.

@tf.function
def for_in_tfrange():
  x = tf.constant(0, dtype=tf.int32)
  for i in tf.range(5):
    x += i
  return x

test_dynamically_unrolled(for_in_tfrange)

for_in_tfrange() uses tf.while_loop.

@tf.function
def for_in_tfdataset():
  x = tf.constant(0, dtype=tf.int64)
  for i in tf.data.Dataset.range(5):
    x += i
  return x

test_dynamically_unrolled(for_in_tfdataset)

for_in_tfdataset() uses tf.data.Dataset.reduce.

@tf.function
def while_py_cond():
  x = 5
  while x > 0:
    x -= 1
  return x

test_dynamically_unrolled(while_py_cond)

while_py_cond() gets unrolled.

@tf.function
def while_tf_cond():
  x = tf.constant(5)
  while x > 0:
    x -= 1
  return x


test_dynamically_unrolled(while_tf_cond)

while_tf_cond() uses tf.while_loop.

tensor に依拠する break か早期の return 節を持つのであれば、top-level 条件や iterable もまた tensor であるべきです。

次のサンプルを比較してください :

@tf.function
def while_py_true_py_break(x):
  while True:  # py true
    if x == 0: # py break
      break
    x -= 1
  return x

test_dynamically_unrolled(while_py_true_py_break, 5)

while_py_true_py_break(5) gets unrolled.

@tf.function
def buggy_while_py_true_tf_break(x):
  while True:   # py true
    if tf.equal(x, 0): # tf break
      break
    x -= 1
  return x

with assert_raises(TypeError):
  test_dynamically_unrolled(buggy_while_py_true_tf_break, 5)

Caught expected exception 
  <class 'TypeError'>: in converted code:

    <ipython-input-34-453240ea98e6>:3 buggy_while_py_true_tf_break  *
        while True:   # py true
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:730 while_stmt
        return _py_while_stmt(test, body, get_state, set_state, init_vars, opts)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:845 _py_while_stmt
        while test(*loop_vars):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:765 __bool__
        self._disallow_bool_casting()
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:531 _disallow_bool_casting
        "using a `tf.Tensor` as a Python `bool`")
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:518 _disallow_when_autograph_enabled
        " decorating it directly with @tf.function.".format(task))

    OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did not convert this function. Try decorating it directly with @tf.function.

@tf.function
def while_tf_true_tf_break(x):
  while tf.constant(True): # tf true
    if x == 0:  # py break
      break
    x -= 1
  return x

test_dynamically_unrolled(while_tf_true_tf_break, 5)

while_tf_true_tf_break(5) uses tf.while_loop.

@tf.function
def buggy_py_for_tf_break():
  x = 0
  for i in range(5):  # py for
    if tf.equal(i, 3): # tf break
      break
    x += i
  return x

with assert_raises(TypeError):
  test_dynamically_unrolled(buggy_py_for_tf_break)

Caught expected exception 
  <class 'TypeError'>: in converted code:

    <ipython-input-36-82742b0a14d0>:4 buggy_py_for_tf_break  *
        for i in range(5):  # py for
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:339 for_stmt
        return _py_for_stmt(iter_, extra_test, body, get_state, set_state, init_vars)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:348 _py_for_stmt
        if extra_test is not None and not extra_test(*state):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:765 __bool__
        self._disallow_bool_casting()
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:531 _disallow_bool_casting
        "using a `tf.Tensor` as a Python `bool`")
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/ops.py:518 _disallow_when_autograph_enabled
        " decorating it directly with @tf.function.".format(task))

    OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did not convert this function. Try decorating it directly with @tf.function.

@tf.function
def tf_for_py_break():
  x = 0
  for i in tf.range(5): # tf for
    if i == 3:  # py break
      break
    x += i
  return x

test_dynamically_unrolled(tf_for_py_break)

tf_for_py_break() uses tf.while_loop.

動的に展開されたループからの結果を累積するために、tf.TensorArray を使用することを望むでしょう。

batch_size = 2
seq_len = 3
feature_size = 4

def rnn_step(inp, state):
  return inp + state

@tf.function
def dynamic_rnn(rnn_step, input_data, initial_state):
  # [batch, time, features] -> [time, batch, features]
  input_data = tf.transpose(input_data, [1, 0, 2])
  max_seq_len = input_data.shape[0]

  states = tf.TensorArray(tf.float32, size=max_seq_len)
  state = initial_state
  for i in tf.range(max_seq_len):
    state = rnn_step(input_data[i], state)
    states = states.write(i, state)
  return tf.transpose(states.stack(), [1, 0, 2])
  
dynamic_rnn(rnn_step,
            tf.random.uniform([batch_size, seq_len, feature_size]),
            tf.zeros([batch_size, feature_size]))

<tf.Tensor: id=1254, shape=(2, 3, 4), dtype=float32, numpy=
array([[[0.20169294, 0.61629033, 0.1433531 , 0.4829756 ],
        [1.1323476 , 1.0190183 , 0.8961072 , 0.9289988 ],
        [1.8864099 , 1.3378067 , 1.4223019 , 1.3112907 ]],

       [[0.6474533 , 0.06091189, 0.8633839 , 0.65446174],
        [0.9992776 , 1.0039562 , 1.1452049 , 1.5506929 ],
        [1.4500049 , 1.0567949 , 1.5264317 , 2.1867056 ]]], dtype=float32)>

tf.cond と同様に、tf.while_loop もまた微妙な点を伴います。- ループは 0 回の実行が可能ですから、while_loop の downstream で使用される総ての tensor は上のループで初期化されなければなりません。- 総てのループ変数の shape/dtype は各 iteration と一環していなければなりません。

@tf.function
def buggy_loop_var_uninitialized():
  for i in tf.range(3):
    x = i
  return x

with assert_raises(ValueError):
  buggy_loop_var_uninitialized()

Caught expected exception 
  <class 'ValueError'>: in converted code:

    <ipython-input-39-815fd6bba8cc>:3 buggy_loop_var_uninitialized  *
        for i in tf.range(3):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:315 for_stmt
        composite_symbol_names)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:419 _tf_range_for_stmt
        _disallow_undefs_into_loop(*init_vars)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:97 _disallow_undefs_into_loop
        ' before the loop: {}'.format(tuple(s.symbol_name for s in undefined)))

    ValueError: TensorFlow requires that the following symbols must be defined before the loop: ('x',)

@tf.function
def f():
  x = tf.constant(0)
  for i in tf.range(3):
    x = i
  return x

f()

<tf.Tensor: id=1309, shape=(), dtype=int32, numpy=2>

@tf.function
def buggy_loop_type_changes():
  x = tf.constant(0, dtype=tf.float32)
  for i in tf.range(3): # Yields tensors of type tf.int32...
    x = i
  return x

with assert_raises(tf.errors.InvalidArgumentError):
  buggy_loop_type_changes()

Got unexpected exception 
  <class 'TypeError'>: in converted code:

    <ipython-input-41-46359fb065eb>:4 buggy_loop_type_changes  *
        for i in tf.range(3): # Yields tensors of type tf.int32...
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:315 for_stmt
        composite_symbol_names)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:478 _tf_range_for_stmt
        opts=opts,
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:769 _tf_while_stmt
        aug_init_vars, **opts)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/control_flow_ops.py:2675 while_loop
        back_prop=back_prop)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/while_v2.py:198 while_loop
        add_control_dependencies=add_control_dependencies)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/func_graph.py:915 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/while_v2.py:176 wrapped_body
        outputs = body(*_pack_sequence_as(orig_loop_vars, args))
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:759 aug_body
        composite_symbol_names)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:195 _verify_tf_loop_vars
        first_iter_var)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/util/nest.py:535 map_structure
        structure[0], [func(*x) for x in entries],
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/util/nest.py:535 
        structure[0], [func(*x) for x in entries],
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:179 _check_same_type
        first_iter_var.dtype.name,

    TypeError: "x" has dtype float32 before the loop, but dtype int32 after one iteration. TensorFlow control flow requires it stays the same.

@tf.function
def buggy_concat():
  x = tf.ones([0, 10])
  for i in tf.range(5):
    x = tf.concat([x, tf.ones([1, 10])], axis=0)
  return x

with assert_raises(ValueError):
  buggy_concat()

Caught expected exception 
  <class 'ValueError'>: in converted code:

    <ipython-input-42-bae298a1ce41>:4 buggy_concat  *
        for i in tf.range(5):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:315 for_stmt
        composite_symbol_names)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:478 _tf_range_for_stmt
        opts=opts,
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:769 _tf_while_stmt
        aug_init_vars, **opts)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/control_flow_ops.py:2675 while_loop
        back_prop=back_prop)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/while_v2.py:198 while_loop
        add_control_dependencies=add_control_dependencies)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/framework/func_graph.py:915 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/ops/while_v2.py:176 wrapped_body
        outputs = body(*_pack_sequence_as(orig_loop_vars, args))
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:759 aug_body
        composite_symbol_names)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:195 _verify_tf_loop_vars
        first_iter_var)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/util/nest.py:535 map_structure
        structure[0], [func(*x) for x in entries],
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/util/nest.py:535 
        structure[0], [func(*x) for x in entries],
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow_core/python/autograph/operators/control_flow.py:191 _check_same_type
        first_iter_shape))

    ValueError: "x" has shape (0, 10) before the loop, but shape (1, 10) after one iteration. TensorFlow control flow requires it stays the same or be more specific.

@tf.function
def concat_with_padding():
  x = tf.zeros([5, 10])
  for i in tf.range(5):
    x = tf.concat([x[:i], tf.ones([1, 10]), tf.zeros([4-i, 10])], axis=0)
    x.set_shape([5, 10])
  return x

concat_with_padding()

<tf.Tensor: id=1432, shape=(5, 10), dtype=float32, numpy=
array([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]], dtype=float32)>

以上