![]() |
![]() |
![]() |
![]() |
TensorFlow 2 の Eager execution はデフォルトで有効になっています。ユーザーインターフェースは直感的で柔軟性に優れていますが(一度限りの演算の実行ははるかに簡単で高速に行われます)、パフォーマンスとデプロイ能力に影響がでることがあります。
TensorFlow 2.0 では Eager Execution の使いやすさとTensorFlow 1.0 のパワーとを同時に提供します。この統合の中核となるのは tf.function
です。これは Python の構文のサブセットを移植可能でハイパフォーマンスな TensorFlow のグラフに変換します。
このチュートリアルでは tf.function
と AutoGraph の基本的な特徴についてひととおり確認します。
主に次の内容と推奨事項について説明しています。
- Eager モードでデバッグしてから、
@tf.function
でデコレートする。 - オブジェクトミューテーションまたはリストの追加といった Python 側の効果に依存しないこと。
tf.function
は TensorFlow 演算子と最も相性が良く、NumPy と Python 呼び出しは定数に変換される。
セットアップ
import tensorflow as tf
発生する可能性のあるエラーの種類を示すヘルパー関数を定義します。
import traceback
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))
traceback.print_exc(limit=2)
except Exception as e:
raise e
else:
raise Exception('Expected {} to be raised but no error was raised!'.format(
error_class))
基礎
使い方
定義する Function
(@tf.function
デコレーターを適用するなどして)は、コアの TensorFlow 演算とまったく変わりません。Eager での実行や勾配の計算などを行えます。
@tf.function # The decorator converts `add` into a `Function`.
def add(a, b):
return a + b
add(tf.ones([2, 2]), tf.ones([2, 2])) # [[2., 2.], [2., 2.]]
2021-08-14 01:20:13.818511: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:13.827594: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:13.828644: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:13.830474: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 AVX512F FMA To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags. 2021-08-14 01:20:13.831129: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:13.832147: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:13.833094: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:14.483952: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:14.484967: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:14.485890: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:937] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero 2021-08-14 01:20:14.486877: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1510] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 14648 MB memory: -> device: 0, name: Tesla V100-SXM2-16GB, pci bus id: 0000:00:05.0, compute capability: 7.0 2021-08-14 01:20:14.838965: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:185] None of the MLIR Optimization Passes are enabled (registered 2) <tf.Tensor: shape=(2, 2), dtype=float32, numpy= array([[2., 2.], [2., 2.]], dtype=float32)>
v = tf.Variable(1.0)
with tf.GradientTape() as tape:
result = add(v, 1.0)
tape.gradient(result, v)
<tf.Tensor: shape=(), dtype=float32, numpy=1.0>
Function
をほかの Function
内で使用できます。
@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: shape=(3, 2), dtype=float32, numpy= array([[3., 3.], [3., 3.], [3., 3.]], dtype=float32)>
Function
は、特に小さな演算が多数含まれるグラフでは、Eager コードよりも高速に実行されることがありますが、高価な演算がいくつか含まれるグラフ(畳み込みなど)では、速度の差はあまり見られません。
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")
2021-08-14 01:20:16.078659: I tensorflow/stream_executor/cuda/cuda_dnn.cc:369] Loaded cuDNN version 8100 Eager conv: 0.004265683000085119 Function conv: 0.004870573000061995 Note how there's not much difference in performance for convolutions 2021-08-14 01:20:16.694021: I tensorflow/core/platform/default/subprocess.cc:304] Start cannot spawn child process: No such file or directory
トレーシング
このセクションは、Function
の内部動作や実装の詳細を説明します。将来的に変更する可能性がありますが、いつなぜトレーシングが発生するのかを理解しておけば、tf.function
を効果的に使用しやすくなります。
「トレーシング」とは?
Function
は TensorFlow Graph でプログラムを実行しますが、tf.Graph
は、Eager TensorFlow プログラムにユーザーが記述するすべてのものを表現することはできません。たとえば、Python はポリモーフィズムをサポートしていますが、tf.Graph
では、その入力に特定のデータ型と次元が必要です。またはコマンドラインの引数を読み取る、エラーを発生させる、より複雑な Python オブジェクトを扱うといったサイドタスクを実施しようとしても、どれも tf.Graph
で実行することはできません。
Function
はコードを 2 つの段階に分けることで、このギャップの橋渡しの役割を果たします。
「トレーシング」と呼ばれる第 1 段階において、
Function
は新しいtf.Graph
を作成します。Python コードは通常通り実行しますが、すべての TensorFlow 演算(2 つのテンソルを加算するなど)は 据え置きとなります。これらはtf.Graph
にとらわれるため、実行しません。第 2 段階では、最初の段階で据え置きとなったすべての演算を含む
tf.Graph
が実行されます。この段階は、トレーシングの段階よりもはるかに高速に行われます。
Function
は、その入力によっては必ずしも最初の段階で呼び出されたときに実行するわけではありません。この判定がどのように行われるのかについては、以下の「トレーシングの規則」をご覧ください。最初の段階を省略して 2 番目の段階のみを実行できれば、TensorFlow の高いパフォーマンスが発揮されます。
Function
がトレーシングしないと判断した場合、トレーシング段階の直後に 第 2 段階が始まるため、Function
を呼び出すと、tf.Graph
の作成と実行が行われます。後の方で、get_concrete_function
を使ってトレーシング段階のみを実行する方法を説明します。
型の異なる引数を Function
に渡すと、両方の段階が実行されます。
@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)
同じ型の引数で Function
を繰り返し呼び出すと、生成されるグラフはまったく同じになるため、TensorFlow はトレーシング段階を省略して前にトレーシングしたグラフを再利用することに注意してください。
# This doesn't print 'Tracing with ...'
print(double(tf.constant("b")))
tf.Tensor(b'bb', shape=(), dtype=string)
すべての利用可能なトレースを確認するには、pretty_printed_concrete_signatures()
を使用できます。
print(double.pretty_printed_concrete_signatures())
double(a) Args: a: int32 Tensor, shape=() Returns: int32 Tensor, shape=() double(a) Args: a: float32 Tensor, shape=() Returns: float32 Tensor, shape=() double(a) Args: a: string Tensor, shape=() Returns: string Tensor, shape=()
ここまで、tf.function
が TensorFlow のグラフトレーシングロジックにキャッシュされた動的ディスパッチレイヤーを作成するのを見てきました。用語についてより具体的に説明すると、次のように言えます。
tf.Graph
は、言語に依存しない、生の移植可能な TensorFlow 計算の表現です。ConcreteFunction
はtf.Graph
をラップします。Function
はConcreteFunction
のキャッシュを管理し、入力に適したものを選択します。tf.function
は Python 関数をラップし、Function
オブジェクトを返します。- トレーシングは
tf.Graph
を作成し、それをConcreteFunction
(またはトレース)をラップします。
トレーシングの規則
Function
は、トレーシングされた ConcreteFunction
を再利用するかどうかを判定します。この判定は、入力の引数とキーワードからキャッシュキーを計算して行われます。キャッシュキーは、次の規則(変更される可能性があります)に従って、Function
呼び出しの入力引数とキーワードに基づく ConcreteFunction
を識別するキーです。
tf.Tensor
に生成されたキーは、その形状と dtype である。tf.Variable
に生成されたキーは、一意の変数 ID である。- Python プリミティブ型(
int
、float
、str
など)に生成されたキーはその値である。 - ネストされた
dict
、list
、tuple
、namedtuple
、およびattr
に生成されたキーは、フラット化されたリーフキーのタプルである(nest.flatten
を参照)。(このフラット化の結果、ネスト構造がトレーシング時とは異なる具象関数が呼び出されると、TypeError が発生します。) - そのたすべての Python の型については、キーはオブジェクト固有です。このため、関数またはメソッドは呼び出しに使用されるインスタンスごとにトレースされます。
注意: キャッシュキーは、Function
の入力パラメータに基づくため、グローバルと自由変数を変更するだけでは、新しいトレースは作成されません。Python のグローバル変数と自由変数を扱う際の推奨される方法については、こちらのセクションをご覧ください。
リトレーシングの制御
リトレーシングは、Function
が 2 つ以上のトレースを作成する際に発生します。これは、TensorFlow が一連の入力ごとに正しいグラフを生成する上で役立ちますが、トレーシングは高価な演算です。Function
が呼び出しごとに新しいグラフをリトレーシングすると、コードの実行は tf.function
を使用しない場合よりも遅くなってしまいます。
トレーシングの動作を制御するには、次のテクニックを使用できます。
- トレーシングを制限するために、
input_signature
をtf.function
に指定します。
@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]]))
# We specified an int32 dtype in the input signature, so this should fail.
with assert_raises(ValueError):
next_collatz(tf.constant([1.0, 2.0]))
Tracing with Tensor("x:0", shape=(None,), dtype=int32) tf.Tensor([4 1], shape=(2,), dtype=int32) Caught expected exception <class 'ValueError'>: Caught expected exception <class 'ValueError'>: Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/3215754498.py", line 9, in <module> next_collatz(tf.constant([[1, 2], [3, 4]])) 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)) Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/3215754498.py", line 13, in <module> next_collatz(tf.constant([1.0, 2.0])) ValueError: Python inputs incompatible with input_signature: inputs: ( tf.Tensor([1. 2.], shape=(2,), dtype=float32)) input_signature: ( TensorSpec(shape=(None,), dtype=tf.int32, name=None))
トレースを柔軟に再利用できるようにするために、[None] 次元を
tf.TensorSpec
に指定します。TensorFlow は形状に基づいてテンソルを一致させるため、ワイルドカードとして
None
次元を使用することで、Function
が可変サイズの入力にトレースを再利用できるようになります。可変サイズの入力は、長さの異なるシーケンスがある場合や、バッチごとに画像のサイズが異なる場合に発生します(例として、Transformer と Deep Dream チュートリアルをご覧ください)。
@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def g(x):
print('Tracing with', x)
return x
# No retrace!
print(g(tf.constant([1, 2, 3])))
print(g(tf.constant([1, 2, 3, 4, 5])))
Tracing with Tensor("x:0", shape=(None,), dtype=int32) tf.Tensor([1 2 3], shape=(3,), dtype=int32) tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)
リトレーシングを減らすために、テンソルに Python 引数をキャストします。
通常、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 = ", num_steps)
tf.print("Executing with num_steps = ", num_steps)
for _ in tf.range(num_steps):
train_one_step()
print("Retracing occurs for different Python arguments.")
train(num_steps=10)
train(num_steps=20)
print()
print("Traces are reused for Tensor arguments.")
train(num_steps=tf.constant(10))
train(num_steps=tf.constant(20))
Retracing occurs for different Python arguments. Tracing with num_steps = 10 Executing with num_steps = 10 Tracing with num_steps = 20 Executing with num_steps = 20 Traces are reused for Tensor arguments. Tracing with num_steps = Tensor("num_steps:0", shape=(), dtype=int32) Executing with num_steps = 10 Executing with num_steps = 20
リトレーシングを強制する必要がある場合は、新しい Function
を作成します。トレースは絶対に、各 Function
オブジェクト間で共有されることはありません。
def f():
print('Tracing!')
tf.print('Executing')
tf.function(f)()
tf.function(f)()
Tracing! Executing Tracing! Executing
具象関数の取得
関数がトレースされるたびに新しい具象関数が作成されますが、get_concrete_function
を使うことで、具象関数を直接取得できます。
print("Obtaining concrete trace")
double_strings = double.get_concrete_function(tf.constant("a"))
print("Executing traced function")
print(double_strings(tf.constant("a")))
print(double_strings(a=tf.constant("b")))
Obtaining concrete trace Executing traced function tf.Tensor(b'aa', shape=(), dtype=string) tf.Tensor(b'bb', shape=(), dtype=string)
# You can also call get_concrete_function on an InputSpec
double_strings_from_inputspec = double.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.string))
print(double_strings_from_inputspec(tf.constant("c")))
Tracing with Tensor("a:0", shape=(), dtype=string) tf.Tensor(b'cc', shape=(), dtype=string)
ConcreteFunction
を出力すると、入力引数(型付き)とその出力型の概要が表示されます。
print(double_strings)
ConcreteFunction double(a) Args: a: string Tensor, shape=() Returns: string Tensor, shape=()
また、具象関数のシグネチャを直接取得することもできます。
print(double_strings.structured_input_signature)
print(double_strings.structured_outputs)
((TensorSpec(shape=(), dtype=tf.string, name='a'),), {}) Tensor("Identity:0", shape=(), dtype=string)
互換性のない型で具象トレースを使用すると、エラーが発生します。
with assert_raises(tf.errors.InvalidArgumentError):
double_strings(tf.constant(1))
Caught expected exception <class 'tensorflow.python.framework.errors_impl.InvalidArgumentError'>: Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/3196284684.py", line 2, in <module> double_strings(tf.constant(1)) tensorflow.python.framework.errors_impl.InvalidArgumentError: cannot compute __inference_double_162 as input #0(zero-based) was expected to be a string tensor but is a int32 tensor [Op:__inference_double_162]
Python 引数は、具象関数の入力シグネチャで特別に扱われていることに気づいたかもしれません。TensorFlow 2.3 より前では、Python 引数は単に具象関数のシグネチャから削除されていましたが、TensorFlow 2.3 からはシグネチャに残されたまま、トレーシング中に値セットを取るように制約されています。
@tf.function
def pow(a, b):
return a ** b
square = pow.get_concrete_function(a=tf.TensorSpec(None, tf.float32), b=2)
print(square)
ConcreteFunction pow(a, b=2) Args: a: float32 Tensor, shape=<unknown> Returns: float32 Tensor, shape=<unknown>
assert square(tf.constant(10.0)) == 100
with assert_raises(TypeError):
square(tf.constant(10.0), b=3)
Caught expected exception <class 'TypeError'>: Traceback (most recent call last): File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1721, in _call_impl cancellation_manager) File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1766, in _call_with_flat_signature self._flat_signature_summary(), ", ".join(sorted(kwargs)))) TypeError: pow(a) got unexpected keyword arguments: b. During handling of the above exception, another exception occurred: Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/2310937119.py", line 4, in <module> square(tf.constant(10.0), b=3) TypeError: ConcreteFunction pow(a, b) was constructed with int value 2 in b, but was called with int value 3
グラフの取得
それぞれの具象関数は、tf.Graph
を囲む呼び出し可能なラッパーです。通常、実際の tf.Graph
オブジェクトを取得する必要はないにしろ、具象関数から簡単に取得することが可能です。
graph = double_strings.graph
for node in graph.as_graph_def().node:
print(f'{node.input} -> {node.name}')
[] -> a ['a', 'a'] -> add ['add'] -> Identity
デバッグ
一般的に、コードのデバックは、tf.function
内で行うよりも、Eager モードで行う方が簡単です。Eager モードでは、tf.function
でデコレートする前に、コードがエラーなく実行することを確認しておく必要があります。デバッグプロセスを支援する目的で、tf.config.run_functions_eagerly(True)
を呼び出すと、tf.function
をグローバルに無効にして、有効にし直すことができます。
tf.function
内でのみ出現する問題を追跡する場合、次のようなヒントがあります。
- 従来のシンプルな Python
print
呼び出しは、トレーシング中にのみ実行されるため、関数が(リ)トレーシングされるときに追跡しやすくなります。 tf.print
呼び出しは毎回実行するため、実行中の中間値の追跡に役立ちます。tf.debugging.enable_check_numerics
は、NaN と Inf がいつ作成されるかを簡単に追跡できます。pdb
は、トレーシング中に何が起きているのかを理解する上で役立ちます。(注意: PDB が示すのは、AutoGraph 変換ソースコードです。)
AutoGraph 変換
AutoGraph は、tf.function
内でデフォルトで利用できるようになっているライブラリで、Python の Eager コードのサブセットとグラフ対応の TensorFlow 演算に変換します。これには、if
、for
、while
などの制御フローが含まれます。
tf.cond
や tf.while_loop
などの TensorFlow 演算は機能し続けますが、制御フローは、Python で記述された場合の方が書きやすく理解しやすいことがほとんどです。
# 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.458563447 0.565437675 0.916076064 0.678965449 0.30448842] [0.428912699 0.512000859 0.7240358 0.590846419 0.295414716] [0.40441224 0.471502692 0.619402945 0.530504107 0.287110865] [0.383717924 0.439412653 0.550712168 0.485766321 0.279473454] [0.365931898 0.413157493 0.501053751 0.450849771 0.272417665] [0.350428253 0.391150385 0.462945491 0.422597259 0.265873015] [0.336755276 0.372351527 0.432481796 0.399116218 0.259780467] [0.32457754 0.356046855 0.407393336 0.379192531 0.25409013] [0.313640058 0.341727227 0.386257112 0.362006 0.248759553] [0.303744942 0.329018503 0.368129015 0.346979707 0.243752241] [0.294736 0.317638576 0.352354079 0.333694249 0.239036724] [0.286487937 0.307370126 0.338461578 0.321836293 0.234585658] [0.278899103 0.298042715 0.32610321 0.311166376 0.230375171] [0.271885842 0.289520383 0.315015 0.301497817 0.226384312] [0.265378714 0.28169331 0.304992527 0.292682737 0.222594574] [0.259319514 0.274471551 0.295874774 0.284602106 0.218989611] [0.253658921 0.267780662 0.287532926 0.277159035 0.215554819] [0.248355 0.261558503 0.279862553 0.270273656 0.212277189] [0.24337168 0.255752683 0.272777855 0.263879448 0.209145114] [0.238677874 0.250318587 0.266207725 0.257920474 0.206148058] [0.234246537 0.245218083 0.260092586 0.252349436 0.203276619] [0.230054021 0.240418345 0.254382104 0.247125864 0.200522214] [0.226079613 0.23589085 0.249033451 0.242215022 0.197877109] [0.22230497 0.23161073 0.244009867 0.237586811 0.195334256] [0.218713865 0.227556229 0.239279598 0.233215094 0.192887247] [0.215291873 0.223708153 0.234815165 0.229076952 0.190530151] [0.212026089 0.2200495 0.230592504 0.225152254 0.18825759] [0.208904967 0.216565207 0.226590499 0.221423253 0.186064631] [0.205918133 0.21324186 0.222790554 0.217874154 0.183946759] [0.203056186 0.210067376 0.219176158 0.214490935 0.181899741] [0.200310647 0.207030967 0.215732694 0.211261019 0.179919735] [0.197673827 0.204122886 0.212447032 0.208173171 0.178003147] <tf.Tensor: shape=(5,), dtype=float32, numpy= array([0.19513872, 0.20133433, 0.2093075 , 0.20521726, 0.17614664], dtype=float32)>
興味があれば、AutoGraph が生成するコードを検査できます。
print(tf.autograph.to_code(f.python_function))
def tf__f(x): with ag__.FunctionScope('f', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope: do_return = False retval_ = ag__.UndefinedReturnValue() def get_state(): return (x,) def set_state(vars_): nonlocal x (x,) = vars_ def loop_body(): nonlocal x ag__.converted_call(ag__.ld(tf).print, (ag__.ld(x),), None, fscope) x = ag__.converted_call(ag__.ld(tf).tanh, (ag__.ld(x),), None, fscope) def loop_test(): return (ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) > 1) ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('x',), {}) try: do_return = True retval_ = ag__.ld(x) except: do_return = False raise return fscope.ret(retval_, do_return)
条件文
AutoGraph は if <condition>
文を相当する tf.cond
呼び出しに変換します。この置換は、<condition>
がテンソルである場合に行われます。テンソルでない場合は、if
文は Python の条件文として実行されます。
Python 条件文はトレーシング中に実行するため、条件文のブランチが 1 つだけグラフに追加されます。AutoGraph を使用しない場合、データに依存する制御フローが存在すると、トレーシングされたこのグラフは別のブランチを取ることができません。
tf.cond
は、条件文の両方のブランチをトレーシングし、実行時に動的に 1 つのブランチを選択してグラフに追加します。トレーシングには意図しない副作用がある場合があります。詳細は、AutoGraph のトレーシング効果をご覧ください。
@tf.function
def fizzbuzz(n):
for i in tf.range(1, n + 1):
print('Tracing for loop')
if i % 15 == 0:
print('Tracing fizzbuzz branch')
tf.print('fizzbuzz')
elif i % 3 == 0:
print('Tracing fizz branch')
tf.print('fizz')
elif i % 5 == 0:
print('Tracing buzz branch')
tf.print('buzz')
else:
print('Tracing default branch')
tf.print(i)
fizzbuzz(tf.constant(5))
fizzbuzz(tf.constant(20))
Tracing for loop Tracing fizzbuzz branch Tracing fizz branch Tracing buzz branch Tracing default branch 1 2 fizz 4 buzz 1 2 fizz 4 buzz fizz 7 8 fizz buzz 11 fizz 13 14 fizzbuzz 16 17 fizz 19 buzz
AutoGraph 変換の if 文におけるその他の制約事項については、リファレンスドキュメントをご覧ください。
ループ
AutoGraph は、一部の for
文と while
文を相当する tf.while_loop
などの TensorFlow のループ演算に変換します。変換されない場合、for
または while
ループは Python ループとして実行されます。
この置き換えは、次の場合に行われます。
for x in y
:y
がテンソルである場合、tf.while_loop
に変換されます。y
がtf.data.Dataset
である特別なケースでは、tf.data.Dataset
演算の組み合わせが生成されます。while <condition>
:<condition>
がテンソルである場合、tf.while_loop
に変換されます。
Python ループは、トレーシング中に実行され、ループの反復ごとに、tf.Graph
に追加の演算が追加されます。
TensorFlow ループはループの本体をトレーシングし、実行時に実行する反復回数を動的に選択します。ループ本体は、生成された tf.Graph
に一度だけ出現します。
AutoGraph 変換の for
文と while
文におけるその他の制約事項については、リファレンスドキュメントをご覧ください。
Python データのループ
一般的な落とし穴は、tf.function
内で Python/Numpy データをループする際にあります。このループは、トレーシングプロセス中に実行し、ループの反復ごとに
モデルのコピーを tf.Graph
に追加してしまいます。
トレーニングループ全体を tf.function
にラップしたいのであれば、データを tf.data.Dataset
としてラップし、AutoGraph にトレーニングループを動的に展開させるようにするのが最も安全な方法です。
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)] * 3
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), (1, 1)]) contains 11 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(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 6 nodes in its graph train(<FlatMapDataset shapes: (<unknown>, <unknown>), types: (tf.int32, tf.int32)>) contains 6 nodes in its graph
Python/Numpy データを Dataset にラップする際は、tf.data.Dataset.from_generator
と tf.data.Dataset.from_tensors
の違いに注意してください。前者は、データを Python に維持し、tf.py_function
経由で取得するため、パフォーマンスに問題がありますが、後者は、データのコピーをグラフ内の大型の tf.constant()
ノードとしてバンドル化するため、メモリに問題が現れます。
データを消費するには、TFRecordDataset/CsvDataset などでファイルからデータを読み取るのが最も効果的な方法です。そうすれば、Python を使わずに、TensorFlow 自体でデータの非同期読み込みとプリフェッチを管理できるようになります。詳細は、tf.data guide をご覧ください。
ループでの値の累積
ループの反復ごとに値を累積していくのは一般的なパターンです。通常は、Python のリストに追加したり、Python ディレクトリにエントリを追加したりして行われますが、これらは Python の副作用であるため、動的に展開されるループでは期待どおりに動作しません。動的に展開されるループの結果を累積する場合は、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: shape=(2, 3, 4), dtype=float32, numpy= array([[[0.35414898, 0.05380535, 0.42085373, 0.06279671], [1.260183 , 0.61282504, 0.6820253 , 1.0446292 ], [1.8472321 , 1.3193153 , 1.0124825 , 1.0583038 ]], [[0.4953791 , 0.4930824 , 0.7954943 , 0.10160744], [1.28636 , 1.3122512 , 0.89200366, 0.8784492 ], [2.0779896 , 1.949412 , 1.8605278 , 1.651388 ]]], dtype=float32)>
制限事項
TensorFlow の Function
には、設計上、いくつかの制限事項があり、Python 関数を Function
に変換する際には、注意が必要です。
Python の副作用の実行
Function
内での出力、リストへのアペンド、グローバル変数のミューテーションといった副作用は、2 回実行されたり、まったく実行しなかったりといったように、予測のつかない動作をすることがあります。また、入力セットで Function
を初めて呼び出した場合にのみ実行し、以降では、Python コードを実行せずに、トレーシング済みの tf.Graph
が再実行されてしまうこともあります。
基本的に、ロジックでは Python の副作用に依存しないようにし、トレースをデバッグするためだけに使用することをお勧めします。呼び出しごとに TensorFlow ランタイムが確実にコードを実行できるようにするには、tf.data
、tf.print
、tf.summary
、tf.Variable.assign
、tf.TensorArray
などの TensorFlow API を使用するのが最善の方法です。
@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
Function
の呼び出しごとに Python コードを実行する場合は、tf.py_function
が脱出口です。code2}tf.py_function には移植性がなく、特にパフォーマンスに優れているわけでもなく、SavedModel で保存できなければ、分散型(マルチ GPU、TPU)の環境でうまく動作するわけでもありません。また、tf.py_function
はグラフに組み込む必要もあるため、すべての入力/出力をテンソルにキャストしてしまいます。
Python のグローバル変数と自由変数の変更
Python のグローバル変数と自由変数の変更は、Python の副作用としてみなされるため、トレーシング中にのみ発生します。
external_list = []
@tf.function
def side_effect(x):
print('Python side effect')
external_list.append(x)
side_effect(1)
side_effect(1)
side_effect(1)
# The list append only happened once!
assert len(external_list) == 1
Python side effect
リスト、辞書、Function
の外側で機能するその他のオブジェクトなどのコンテナのミューテーションは避けてください。代わりに、引数と TF オブジェクトを使用しましょう。たとえば、「ループでの値の累積」セクションには、リストのような演算を実装する方法の一例が示されています。
一部のケースでは、tf.Variable
である場合に状態をキャプチャして操作することができます。Keras モデルの重みは、このようにして、同じ ConcreteFunction
への呼び出しの繰り返しで更新されています。
Python イテレータとジェネレータの使用
ジェネレータやイテレータなどの多くの Python 機能は、Python ランタイムに依存して状態を追跡しています。一般的に、これらのコンストラクトは Eager モードでも期待どおりに動作しますが、Python の副作用の例であるため、トレーシング中にしか発生しません。
@tf.function
def buggy_consume_next(iterator):
tf.print("Value:", next(iterator))
iterator = iter([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: 1 Value: 1 Value: 1
TensorFlow にリストコントラクト用の特別な tf.TensorArray
があるように、イテレーション用にも特別な tf.data.Iterator
があります。概要は、AutoGraph 変換をご覧ください。また、tf.data
API を使って、ジェネレータのパターンを実装できます。
@tf.function
def good_consume_next(iterator):
# This is ok, iterator is a tf.data.Iterator
tf.print("Value:", next(iterator))
ds = tf.data.Dataset.from_tensor_slices([1, 2, 3])
iterator = iter(ds)
good_consume_next(iterator)
good_consume_next(iterator)
good_consume_next(iterator)
Value: 1 Value: 2 Value: 3
Function
呼び出し間での tf.Variables の削除
発生する可能性のあるもう 1 つのエラーは、ガベージコレクションされた変数です。ConcreteFunction
は使用した変数への WeakRefs のみを保持するため、ユーザーが変数への参照を保持するようにする必要があります。
external_var = tf.Variable(3)
@tf.function
def f(x):
return x * external_var
traced_f = f.get_concrete_function(4)
print("Calling concrete function...")
print(traced_f(4))
# The original variable object gets garbage collected, since there are no more
# references to it.
external_var = tf.Variable(4)
print()
print("Calling concrete function after garbage collecting its closed Variable...")
with assert_raises(tf.errors.FailedPreconditionError):
traced_f(4)
Calling concrete function... tf.Tensor(12, shape=(), dtype=int32) Calling concrete function after garbage collecting its closed Variable... Caught expected exception <class 'tensorflow.python.framework.errors_impl.FailedPreconditionError'>: Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/3862898592.py", line 16, in <module> traced_f(4) tensorflow.python.framework.errors_impl.FailedPreconditionError: 2 root error(s) found. (0) Failed precondition: Could not find variable _AnonymousVar3. This could mean that the variable has been deleted. In TF1, it can also mean the variable is uninitialized. Debug info: container=localhost, status=Not found: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist. [[node ReadVariableOp (defined at tmp/ipykernel_11770/3862898592.py:4) ]] [[ReadVariableOp/_2]] (1) Failed precondition: Could not find variable _AnonymousVar3. This could mean that the variable has been deleted. In TF1, it can also mean the variable is uninitialized. Debug info: container=localhost, status=Not found: Resource localhost/_AnonymousVar3/N10tensorflow3VarE does not exist. [[node ReadVariableOp (defined at tmp/ipykernel_11770/3862898592.py:4) ]] 0 successful operations. 0 derived errors ignored. [Op:__inference_f_772] Function call stack: f -> f
既知の問題
Function
が正しく評価していない場合、以下の既知の問題が該当する可能性があります。これらの問題は、今後修正される予定です。
Python のグローバル変数と自由変数への依存
Function
は、Python 引数の新しい値で呼び出された時に新しい ConcreteFunction
を作成しますが、Python クロージャ、グローバル変数、またはその Function
の非ローカル変数に対しては作成しません。Function
への呼び出しごとに値が変化する場合でも、Function
はトレーシングされたときの値をそのまま使用してしまいます。これは、通常の Python 関数の動作とは異なります。
このため、外側の名前を閉じる代わりに引数を使用する関数プログラミングの様式をお勧めします。
@tf.function
def buggy_add():
return 1 + foo
@tf.function
def recommended_add(foo):
return 1 + foo
foo = 1
print("Buggy:", buggy_add())
print("Correct:", recommended_add(foo))
Buggy: tf.Tensor(2, shape=(), dtype=int32) Correct: tf.Tensor(2, shape=(), dtype=int32)
print("Updating the value of `foo` to 100!")
foo = 100
print("Buggy:", buggy_add()) # Did not change!
print("Correct:", recommended_add(foo))
Updating the value of `foo` to 100! Buggy: tf.Tensor(2, shape=(), dtype=int32) Correct: tf.Tensor(101, shape=(), dtype=int32)
外側の名前は、その値を更新しない場合にのみ閉じることができます。
Python オブジェクトへの依存
Python オブジェクトを引数として tf.function
に渡す上での推奨事項には、多数の既知の問題があります。これらは今後修正される予定です。一般的に、Python のプリミティブ型または tf.nest
と互換性のある構造を引数として使用する場合や、オブジェクトの別のインスタンスを Function
に渡す場合には、一貫したトレーシングを期待できますが、同一のオブジェクトであっても属性が異なるものを渡す場合、Function
は、新しいトレースを作成しません。
class SimpleModel(tf.Module):
def __init__(self):
# These values are *not* tf.Variables.
self.bias = 0.
self.weight = 2.
@tf.function
def evaluate(model, x):
return model.weight * x + model.bias
simple_model = SimpleModel()
x = tf.constant(10.)
print(evaluate(simple_model, x))
tf.Tensor(20.0, shape=(), dtype=float32)
print("Adding bias!")
simple_model.bias += 5.0
print(evaluate(simple_model, x)) # Didn't change :(
Adding bias! tf.Tensor(20.0, shape=(), dtype=float32)
同じ Function
を使用して、更新されたモデルのインスタンスを評価する場合、更新されたモデルには、元のモデルと同じキャッシュキーが含まれるため、不具合が生じます。
このため、ミュート可能なオブジェクト属性に依存しない Function
を記述するか、新しいオブジェクトを作成することをお勧めします。
この方法が困難な場合は、回避策として、オブジェクトを変更するたびに新しい Function
がリトレーシングを行うようにする方法が挙げられます。
def evaluate(model, x):
return model.weight * x + model.bias
new_model = SimpleModel()
evaluate_no_bias = tf.function(evaluate).get_concrete_function(new_model, x)
# Don't pass in `new_model`, `Function` already captured its state during tracing.
print(evaluate_no_bias(x))
tf.Tensor(20.0, shape=(), dtype=float32)
print("Adding bias!")
new_model.bias += 5.0
# Create new Function and ConcreteFunction since you modified new_model.
evaluate_with_bias = tf.function(evaluate).get_concrete_function(new_model, x)
print(evaluate_with_bias(x)) # Don't pass in `new_model`.
Adding bias! tf.Tensor(25.0, shape=(), dtype=float32)
リトレーシングにはコストがかかるため、tf.Variable
をオブジェクト属性として使用することができます。こうすることで、リトレーシングを行わずに、ミュートして(変更はしません)同様の効果を得ることができます。
class BetterModel:
def __init__(self):
self.bias = tf.Variable(0.)
self.weight = tf.Variable(2.)
@tf.function
def evaluate(model, x):
return model.weight * x + model.bias
better_model = BetterModel()
print(evaluate(better_model, x))
tf.Tensor(20.0, shape=(), dtype=float32)
print("Adding bias!")
better_model.bias.assign_add(5.0) # Note: instead of better_model.bias += 5
print(evaluate(better_model, x)) # This works!
Adding bias! tf.Tensor(25.0, shape=(), dtype=float32)
tf.Variables の作成
Function
では、最初に呼び出された時に 1 回だけ変数を作成し、以降ではそれが再利用されます。新しいトレースで tf.Variables
を作成することはできません。以降の呼び出して新しい変数を作成することはできませんが、将来的には可能になる予定です。
例:
@tf.function
def f(x):
v = tf.Variable(1.0)
return v
with assert_raises(ValueError):
f(1.0)
Caught expected exception <class 'ValueError'>: Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/3018268426.py", line 7, in <module> f(1.0) ValueError: in user code: /tmp/ipykernel_11770/3018268426.py:3 f * v = tf.Variable(1.0) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:268 __call__ ** return cls._variable_v2_call(*args, **kwargs) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:262 _variable_v2_call shape=shape) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter return captured_getter(captured_previous, **kwargs) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:765 invalid_creator_scope "tf.function-decorated function tried to create " ValueError: tf.function-decorated function tried to create variables on non-first call.
関数が初めて実行される際に作成される変数である限り、それらの変数を Function
内で作成できます。
class Count(tf.Module):
def __init__(self):
self.count = None
@tf.function
def __call__(self):
if self.count is None:
self.count = tf.Variable(0)
return self.count.assign_add(1)
c = Count()
print(c())
print(c())
tf.Tensor(1, shape=(), dtype=int32) tf.Tensor(2, shape=(), dtype=int32)
複数の Keras オプティマイザとの使用
2 つ以上の Keras オプティマイザを tf.function
で使用しようとすると、「ValueError: tf.function-decorated function tried to create variables on non-first call.
」というエラーが発生することがあります。このエラーは、オプティマイザが初めて勾配を適用する際に、内部的に tf.Variables
を作成するために発生するものです。
opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)
opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)
@tf.function
def train_step(w, x, y, optimizer):
with tf.GradientTape() as tape:
L = tf.reduce_sum(tf.square(w*x - y))
gradients = tape.gradient(L, [w])
optimizer.apply_gradients(zip(gradients, [w]))
w = tf.Variable(2.)
x = tf.constant([-1.])
y = tf.constant([2.])
train_step(w, x, y, opt1)
print("Calling `train_step` with different optimizer...")
with assert_raises(ValueError):
train_step(w, x, y, opt2)
Calling `train_step` with different optimizer... Caught expected exception <class 'ValueError'>: Traceback (most recent call last): File "/tmp/ipykernel_11770/3551158538.py", line 8, in assert_raises yield File "/tmp/ipykernel_11770/3167358578.py", line 18, in <module> train_step(w, x, y, opt2) ValueError: in user code: /tmp/ipykernel_11770/3167358578.py:9 train_step * optimizer.apply_gradients(zip(gradients, [w])) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:628 apply_gradients ** self._create_all_weights(var_list) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:813 _create_all_weights _ = self.iterations /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:820 __getattribute__ return super(OptimizerV2, self).__getattribute__(name) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:980 iterations aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py:1186 add_weight aggregation=aggregation) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py:818 _add_variable_with_custom_getter **kwargs_for_getter) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/engine/base_layer_utils.py:129 make_variable shape=variable_shape if variable_shape else None) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:266 __call__ return cls._variable_v1_call(*args, **kwargs) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:227 _variable_v1_call shape=shape) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/ops/variables.py:67 getter return captured_getter(captured_previous, **kwargs) /tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py:765 invalid_creator_scope "tf.function-decorated function tried to create " ValueError: tf.function-decorated function tried to create variables on non-first call.
トレーニング中にオプティマイザを変更する必要がある場合は、回避策として、オプティマイザごとに新しい Function
を作成し、ConcreteFunction
を直接呼び出すようにすることができます。
opt1 = tf.keras.optimizers.Adam(learning_rate = 1e-2)
opt2 = tf.keras.optimizers.Adam(learning_rate = 1e-3)
# Not a tf.function.
def train_step(w, x, y, optimizer):
with tf.GradientTape() as tape:
L = tf.reduce_sum(tf.square(w*x - y))
gradients = tape.gradient(L, [w])
optimizer.apply_gradients(zip(gradients, [w]))
w = tf.Variable(2.)
x = tf.constant([-1.])
y = tf.constant([2.])
# Make a new Function and ConcreteFunction for each optimizer.
train_step_1 = tf.function(train_step).get_concrete_function(w, x, y, opt1)
train_step_2 = tf.function(train_step).get_concrete_function(w, x, y, opt2)
for i in range(10):
if i % 2 == 0:
train_step_1(w, x, y) # `opt1` is not used as a parameter.
else:
train_step_2(w, x, y) # `opt2` is not used as a parameter.
複数の Keras モデルとの使用
また、別のモデルインスタンスを同一の Function
に渡す際に、「ValueError: tf.function-decorated function tried to create variables on non-first call.
」というエラーも発生することがあります。
このエラーは、Keras モデル(入力形状が定義されていない)と Keras レイヤーが、初めて呼び出されるときに tf.Variables
を作成するために発生するものです。これらの変数をすでに呼び出された Function
内で初期化しようとしているのでしょう。このエラーを回避するには、モデルをトレーニングする前に、model.build(input_shape)
を呼び出して、すべての重みを初期化するようにしてください。
参考資料
Function
のエクスポートと読み込みの方法については、SavedModel ガイドをご覧ください。トレーシングの後に実行するグラフの最適化については、Grappler ガイドをご覧ください。データパイプラインの最適化方法とモデルのプロファイリングについては、Profiler ガイドをご覧ください。