TensorFlow 2.0 Beta : 上級 Tutorials : カスタマイズ :- tf.function

Posted on by Sales Information

TensorFlow 2.0 Beta : 上級 Tutorials : カスタマイズ :- tf.function (翻訳/解説)

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

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

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、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

!pip install -q tensorflow==2.0.0-beta1
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=40, 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=67, 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 の利用方法の微妙な点の殆どはこの retracing 動作から生じています。

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

# 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)

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

  • 新しい tf.function を作成します。別々の tf.function オブジェクトは trace を共有しないことが保証されます。
  • 特定の trace を得るために get_concrete_function メソッドを使用します。
  • グラフ呼び出し毎に一度だけ trace するために 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_98 as input #0(zero-based) was expected to be a string tensor but is a int32 tensor [Op:__inference_double_98]

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def next_collatz(x):
  print("Tracing with", x)
  return tf.where(tf.equal(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: id=125, shape=(2, 2), dtype=int32, numpy=
array([[1, 2],
       [3, 4]], dtype=int32)>,)), input_signature ((TensorSpec(shape=(None,), dtype=tf.int32, name=None),))

 

いつ retrace するか？

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

 

Python あるいは Tensor args?

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

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

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 副作用は tracing の間だけ発生します。それでは tf.function からどのように副作用を確実に引き起こせるのでしょう？

一般的な経験則は trace をデバッグするために Python 副作用だけを使用することです。さもなければ、 tf.Variable.assign, tf.print と tf.summary のような TensorFlow ops は各呼び出しに伴う TensorFlow ランタイムにより貴方のコードが trace されて実行されることを確かなものにする最善の方法です。一般に 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

WARNING: Logging before flag parsing goes to stderr.
W0629 01:08:56.691862 140487172290304 backprop.py:842] The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int32
W0629 01:08:56.694054 140487155504896 backprop.py:842] The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int32
W0629 01:08:56.695917 140487172290304 backprop.py:842] The dtype of the watched tensor must be floating (e.g. tf.float32), got tf.int32

Python side effect
Python side effect
Python side effect

 

Python 状態に注意する

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

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

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 が多分 trace され、これは巨大グラフ (= 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)))

W0629 01:08:57.038732 140490414335744 deprecation.py:323] From /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/data/ops/dataset_ops.py:505: py_func (from tensorflow.python.ops.script_ops) is deprecated and will be removed in a future version.
Instructions for updating:
tf.py_func is deprecated in TF V2. Instead, there are two
    options available in V2.
    - tf.py_function takes a python function which manipulates tf eager
    tensors instead of numpy arrays. It's easy to convert a tf eager tensor to
    an ndarray (just call tensor.numpy()) but having access to eager tensors
    means <a href="../../../versions/r2.0/api_docs/python/tf/py_function"><code>tf.py_function</code></a>s can use accelerators such as GPUs as well as
    being differentiable using a gradient tape.
    - tf.numpy_function maintains the semantics of the deprecated tf.py_func
    (it is not differentiable, and manipulates numpy arrays). It drops the
    stateful argument making all functions stateful.
    

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(<DatasetV1Adapter shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 4 nodes in its graph
train(<DatasetV1Adapter shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 4 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=460, shape=(), dtype=float32, numpy=10.0>

 

Variables

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

特に、これは各呼び出しで新しい Variable を作成するときに発生します。tracing セマンティクスにより、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)

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/python/ops/variables.py:262 __call__
        return cls._variable_v2_call(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/variables.py:256 _variable_v2_call
        shape=shape)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/variables.py:60 getter
        return captured_getter(captured_previous, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/eager/def_function.py:364 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.604465961 0.235631824 0.292372584 0.0933123827 0.8393749]
[0.540219843 0.231365576 0.284317046 0.093042478 0.685477853]
[0.493154347 0.227323771 0.276895881 0.0927749 0.595069051]
[0.456716418 0.223487303 0.270029694 0.0925096273 0.533531547]
[0.427404225 0.219839349 0.263652474 0.0922466218 0.488075942]
[0.403149694 0.21636492 0.257708609 0.0919858441 0.452688]
[0.382640719 0.213050663 0.252151072 0.0917272717 0.424106032]
[0.364998549 0.209884614 0.2469396 0.0914708674 0.40038386]
[0.349609256 0.206856042 0.242039666 0.091216594 0.380277365]
[0.336028934 0.203955248 0.237421349 0.0909644365 0.362948328]
[0.323927581 0.201173469 0.233058617 0.0907143652 0.347808361]
[0.313053906 0.198502809 0.228928685 0.0904663429 0.334430456]
[0.303212792 0.195936009 0.225011513 0.0902203396 0.322496086]
[0.294249982 0.193466514 0.221289411 0.0899763331 0.311762124]
[0.286041766 0.191088319 0.21774666 0.0897343 0.302039325]
[0.278487593 0.188795939 0.214369327 0.0894942135 0.293177754]
[0.27150473 0.186584324 0.21114485 0.0892560408 0.285056978]
[0.265024424 0.184448823 0.208062038 0.0890197605 0.27757892]
[0.258989 0.182385176 0.205110788 0.0887853503 0.270662814]

<tf.Tensor: id=727, shape=(5,), dtype=float32, numpy=
array([0.25334966, 0.1803894 , 0.202282  , 0.08855279, 0.26424146],
      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()

  def loop_test(x_1):
    return ag__.converted_call('reduce_sum', tf, ag__.ConversionOptions(recursive=True, force_conversion=False, optional_features=(), internal_convert_user_code=True), (x_1,), None) > 1

  def loop_body(x_1):
    ag__.converted_call('print', tf, ag__.ConversionOptions(recursive=True, force_conversion=False, optional_features=(), internal_convert_user_code=True), (x_1,), None)
    x_1 = ag__.converted_call('tanh', tf, ag__.ConversionOptions(recursive=True, force_conversion=False, optional_features=(), internal_convert_user_code=True), (x_1,), None)
    return x_1,
  x, = ag__.while_stmt(loop_test, loop_body, (x,))
  do_return = True
  retval_ = x
  cond = ag__.is_undefined_return(retval_)

  def get_state():
    return ()

  def set_state(_):
    pass

  def if_true():
    retval_ = None
    return retval_

  def if_false():
    return retval_
  retval_ = ag__.if_stmt(cond, if_true, if_false, get_state, set_state)
  return retval_

 

AutoGraph: 条件節

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

この置き換えは条件が Tensor である場合に行われます。さもなければ、条件節は tracing の間に実行されます。

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 は多くの微妙な点を持ちます。- それは条件節の両側を trace し、それから条件に依拠して、ランタイムに適切な分岐を選択することにより動作します。両側の trace は 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=802, 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/python/autograph/operators/control_flow.py:439 if_stmt
        return tf_if_stmt(cond, body, orelse, get_state, set_state)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:456 tf_if_stmt
        outputs, final_state = control_flow_ops.cond(cond, body, orelse)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/util/deprecation.py:507 new_func
        return func(*args, **kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/control_flow_ops.py:1147 cond
        return cond_v2.cond_v2(pred, true_fn, false_fn, name)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/cond_v2.py:86 cond_v2
        op_return_value=pred)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/framework/func_graph.py:716 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:486 wrapper
        outputs = func()
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:512 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

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

@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_range)
test_dynamically_unrolled(for_in_tfrange)
test_dynamically_unrolled(for_in_tfdataset)

for_in_range() gets unrolled.
for_in_tfrange() uses tf.while_loop.
for_in_tfdataset() uses tf.data.Dataset.reduce.

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

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

test_dynamically_unrolled(while_py_cond)
test_dynamically_unrolled(while_tf_cond)

while_py_cond() gets unrolled.
while_tf_cond() uses tf.while_loop.

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

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

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

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

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

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

with assert_raises(TypeError):
  test_dynamically_unrolled(buggy_py_for_tf_break)
test_dynamically_unrolled(tf_for_tf_break)

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

    <ipython-input-30-220fba6e1df8>:3 buggy_while_py_true_tf_break  *
        while True:
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:313 while_stmt
        return _py_while_stmt(test, body, init_state, opts)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:401 _py_while_stmt
        while test(*state):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/framework/ops.py:698 __bool__
        raise TypeError("Using a `tf.Tensor` as a Python `bool` is not allowed. "

    TypeError: Using a `tf.Tensor` as a Python `bool` is not allowed. Use `if t is not None:` instead of `if t:` to test if a tensor is defined, and use TensorFlow ops such as tf.cond to execute subgraphs conditioned on the value of a tensor.

while_tf_true_tf_break(5) uses tf.while_loop.
Caught expected exception 
  <class 'TypeError'>: in converted code:

    <ipython-input-30-220fba6e1df8>:24 buggy_py_for_tf_break  *
        for i in range(5):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:110 for_stmt
        return _py_for_stmt(iter_, extra_test, body, init_state)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:117 _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/python/framework/ops.py:698 __bool__
        raise TypeError("Using a `tf.Tensor` as a Python `bool` is not allowed. "

    TypeError: Using a `tf.Tensor` as a Python `bool` is not allowed. Use `if t is not None:` instead of `if t:` to test if a tensor is defined, and use TensorFlow ops such as tf.cond to execute subgraphs conditioned on the value of a tensor.

tf_for_tf_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=1307, shape=(2, 3, 4), dtype=float32, numpy=
array([[[0.2641269 , 0.32148373, 0.27760386, 0.871022  ],
        [0.36447   , 0.82768893, 1.1970165 , 1.2398282 ],
        [1.1957974 , 1.242557  , 1.2163385 , 1.7522211 ]],

       [[0.11496973, 0.8242916 , 0.5828695 , 0.01092482],
        [0.64337647, 1.7383248 , 1.182523  , 0.94805145],
        [1.2105927 , 1.9715753 , 1.2999045 , 1.6146696 ]]], 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

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

with assert_raises(ValueError):
  buggy_loop_var_uninitialized()
f()

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

    <ipython-input-32-4012abce963a>:3 buggy_loop_var_uninitialized  *
        for i in tf.range(3):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:95 for_stmt
        return _known_len_tf_for_stmt(iter_, extra_test, body, init_state)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:125 _known_len_tf_for_stmt
        _disallow_undefs_into_loop(*init_state)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:50 _disallow_undefs_into_loop
        tuple(s.symbol_name for s in undefined)))

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


<tf.Tensor: id=1371, 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()

Caught expected exception 
  <class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>: Input 1 of node while/merge/_10 was passed int32 from while/next_iteration/_28:0 incompatible with expected float. [Op:__inference_buggy_loop_type_changes_1428]

@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()
  
@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()

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

    <ipython-input-34-d212bdeeeb5e>:4 buggy_concat  *
        for i in tf.range(5):
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:95 for_stmt
        return _known_len_tf_for_stmt(iter_, extra_test, body, init_state)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:156 _known_len_tf_for_stmt
        opts=dict(maximum_iterations=n))
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/autograph/operators/control_flow.py:327 _tf_while_stmt
        retval = control_flow_ops.while_loop(test, body, init_state, **opts)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/control_flow_ops.py:2646 while_loop
        return_same_structure=return_same_structure)
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/while_v2.py:213 while_loop
        len_orig_loop_vars], expand_composites=True))
    /tmpfs/src/tf_docs_env/lib/python3.5/site-packages/tensorflow/python/ops/while_v2.py:869 _check_shapes_compat
        "specify a less-specific shape." % (input_t.name, shape, t.shape))

    ValueError: Input tensor 'ones:0' enters the loop with shape (0, 10), but has shape (1, 10) after one iteration. To allow the shape to vary across iterations, use the `shape_invariants` argument of tf.while_loop to specify a less-specific shape.


<tf.Tensor: id=1549, 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)>
 

以上