tf.function によるパフォーマンスの改善

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード

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 を効果的に使用しやすくなります。

「トレーシング」とは?

FunctionTensorFlow Graph でプログラムを実行しますが、tf.Graph は、Eager TensorFlow プログラムにユーザーが記述するすべてのものを表現することはできません。たとえば、Python はポリモーフィズムをサポートしていますが、tf.Graph では、その入力に特定のデータ型と次元が必要です。またはコマンドラインの引数を読み取る、エラーを発生させる、より複雑な Python オブジェクトを扱うといったサイドタスクを実施しようとしても、どれも tf.Graph で実行することはできません。

Function はコードを 2 つの段階に分けることで、このギャップの橋渡しの役割を果たします。

  1. トレーシング」と呼ばれる第 1 段階において、Function は新しい tf.Graph を作成します。Python コードは通常通り実行しますが、すべての TensorFlow 演算(2 つのテンソルを加算するなど)は 据え置きとなります。これらは tf.Graph にとらわれるため、実行しません。

  2. 第 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 計算の表現です。
  • ConcreteFunctiontf.Graph をラップします。
  • FunctionConcreteFunction のキャッシュを管理し、入力に適したものを選択します。
  • tf.function は Python 関数をラップし、Function オブジェクトを返します。
  • トレーシングtf.Graph を作成し、それを ConcreteFunction(またはトレース)をラップします。

トレーシングの規則

Function は、トレーシングされた ConcreteFunction を再利用するかどうかを判定します。この判定は、入力の引数とキーワードからキャッシュキーを計算して行われます。キャッシュキーは、次の規則(変更される可能性があります)に従って、Function 呼び出しの入力引数とキーワードに基づく ConcreteFunction を識別するキーです。

  • tf.Tensor に生成されたキーは、その形状と dtype である。
  • tf.Variable に生成されたキーは、一意の変数 ID である。
  • Python プリミティブ型(intfloatstr など)に生成されたキーはその値である。
  • ネストされた dictlisttuplenamedtuple、および attr に生成されたキーは、フラット化されたリーフキーのタプルである(nest.flatten を参照)。(このフラット化の結果、ネスト構造がトレーシング時とは異なる具象関数が呼び出されると、TypeError が発生します。)
  • そのたすべての Python の型については、キーはオブジェクト固有です。このため、関数またはメソッドは呼び出しに使用されるインスタンスごとにトレースされます。

注意: キャッシュキーは、Function の入力パラメータに基づくため、グローバルと自由変数を変更するだけでは、新しいトレースは作成されません。Python のグローバル変数と自由変数を扱う際の推奨される方法については、こちらのセクションをご覧ください。

リトレーシングの制御

リトレーシングは、Function が 2 つ以上のトレースを作成する際に発生します。これは、TensorFlow が一連の入力ごとに正しいグラフを生成する上で役立ちますが、トレーシングは高価な演算です。Function が呼び出しごとに新しいグラフをリトレーシングすると、コードの実行は tf.function を使用しない場合よりも遅くなってしまいます。

トレーシングの動作を制御するには、次のテクニックを使用できます。

  • トレーシングを制限するために、input_signaturetf.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 が可変サイズの入力にトレースを再利用できるようになります。可変サイズの入力は、長さの異なるシーケンスがある場合や、バッチごとに画像のサイズが異なる場合に発生します(例として、TransformerDeep 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 演算に変換します。これには、ifforwhile などの制御フローが含まれます。

tf.condtf.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 に変換されます。ytf.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_generatortf.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.datatf.printtf.summarytf.Variable.assigntf.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 ガイドをご覧ください。