今日のローカルTensorFlowEverywhereイベントの出欠確認!

グラフと関数の基礎

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

グラフと tf.function の基礎

このガイドは、TensorFlow の仕組みを説明するために、TensorFlow と Keras 基礎を説明します。今すぐ Keras に取り組みたい方は、Keras のガイド一覧を参照してください。

このガイドでは、グラフ取得のための単純なコード変更、格納と表現、およびモデルの高速化とエクスポートを行うための使用方法について、TensorFlow の中核的な仕組みを説明します。

注意: TensorFlow 1.x のみの知識をお持ちの場合は、このガイドでは、非常に異なるグラフビューが紹介されています。

これは、基礎を概説したガイドです。これらの概念の徹底ガイドについては、tf.function ガイドを参照してください。

グラフとは?

前回の 3 つのガイドでは、TensorFlow の Eager execution について説明しました。これは、TensorFlow 演算が演算ごとにPythonによって実行され、結果を Python に返すことを意味します。Eager TensorFlow は GPU を活用し、変数、テンソル、さらには演算を GPU と TPU に配置することができます。また、デバックも簡単に行えます。

一部のユーザーは、Python から移動する必要はありません。

ただし、TensorFlow を Python で演算ごとに実行すると、ほかの方法では得られない多数の高速化機能が利用できなくなります。Python からテンソルの計算を抽出できる場合は、グラフ にすることができます。

グラフとは、計算のユニットを表す一連の tf.Operation オブジェクトと、演算間を流れるデータのユニットを表す tf.Tensor オブジェクトを含むデータ構造です。 tf.Graph コンテキストで定義されます。これらのグラフはデータ構造であるため、元の Python コードがなくても、保存、実行、および復元することができます。

次は、TensorBoard で視覚化された単純な二層グラフです。

a two-layer tensorflow graph

グラフのメリット

グラフを使用すると、柔軟性が大幅に向上し、モバイルアプリケーション。組み込みデバイス、バックエンドサーバーといった Python インタプリタのない環境でも TensorFlow グラフを使用できます。TensorFlow は、Python からエクスポートされた場合に、保存されるモデルの形式としてグラフを使用します。

また、グラフは最適化を簡単に行えるため、コンパイラは次のような変換を行えます。

  • 計算に定数ノードを畳み込むで、テンソルの値を統計的に推論します(「定数畳み込み」)
  • 独立した計算のサブパートを分離し、スレッドまたはデバイスに分割します。
  • 共通部分式を取り除き、算術演算を単純化します。

これやほかの高速化を実行する Grappler という総合的な最適化システムがあります。

まとめると、グラフは非常に便利なもので、複数のデバイスで、TensorFlow の高速化並列化、および効率化を期待することができます。

ただし、便宜上、Python で機械学習モデル(またはその他の計算)を定義した後、必要となったときに自動的にグラフを作成することをお勧めします。

グラフのトレース

TensorFlow でグラフを作成する方法は、直接呼出しまたはデコレータのいずれかとして tf.function を使用することです。

import tensorflow as tf
import timeit
from datetime import datetime
# Define a Python function
def function_to_get_faster(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# Create a `Function` object that contains a graph
a_function_that_uses_a_graph = tf.function(function_to_get_faster)

# Make some tensors
x1 = tf.constant([[1.0, 2.0]])
y1 = tf.constant([[2.0], [3.0]])
b1 = tf.constant(4.0)

# It just works!
a_function_that_uses_a_graph(x1, y1, b1).numpy()
array([[12.]], dtype=float32)

tf.function 化された関数は、Python コーラブルで、Python 相当と同じように機能します。特定のクラス(python.eager.def_function.Function)を使用しますが、ユーザーにとっては、トレースできないものと同じように動作します。

tf.function は、それが呼び出す Python 関数を再帰的にトレースします。

def inner_function(x, y, b):
  x = tf.matmul(x, y)
  x = x + b
  return x

# Use the decorator
@tf.function
def outer_function(x):
  y = tf.constant([[2.0], [3.0]])
  b = tf.constant(4.0)

  return inner_function(x, y, b)

# Note that the callable will create a graph that
# includes inner_function() as well as outer_function()
outer_function(tf.constant([[1.0, 2.0]])).numpy()
array([[12.]], dtype=float32)

TensorFlow 1.x を使用したことがある場合は、Placeholder または tf.Sesssion をまったく定義する必要がないことに気づくでしょう。

フローの制御と副次的影響

フロー制御とループは、デフォルトで tf.autograph によって TensorFlow に変換されます。Autograph は、ループコンストラクトの標準化、アンロール、および AST マニピュレーションなどのメソッドを組み合わせて使用します。

def my_function(x):
  if tf.reduce_sum(x) <= 1:
    return x * x
  else:
    return x-1

a_function = tf.function(my_function)

print("First branch, with graph:", a_function(tf.constant(1.0)).numpy())
print("Second branch, with graph:", a_function(tf.constant([5.0, 5.0])).numpy())
First branch, with graph: 1.0
Second branch, with graph: [4. 4.]

Autograph 変換を直接呼び出して、Python が TensorFlow 演算に変換される様子を確認することができます。これはほとんど解読不能ですが、変換を確認することができます。

# Don't read the output too carefully.
print(tf.autograph.to_code(my_function))
def tf__my_function(x):
    with ag__.FunctionScope('my_function', '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 (do_return, retval_)

        def set_state(vars_):
            nonlocal do_return, retval_
            (do_return, retval_) = vars_

        def if_body():
            nonlocal do_return, retval_
            try:
                do_return = True
                retval_ = (ag__.ld(x) * ag__.ld(x))
            except:
                do_return = False
                raise

        def else_body():
            nonlocal do_return, retval_
            try:
                do_return = True
                retval_ = (ag__.ld(x) - 1)
            except:
                do_return = False
                raise
        ag__.if_stmt((ag__.converted_call(ag__.ld(tf).reduce_sum, (ag__.ld(x),), None, fscope) <= 1), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
        return fscope.ret(retval_, do_return)


Autograph は、if-then 句、ループ、 breakreturncontinue などを自動的に変換します。

ほとんどの場合、Autograph の動作に特別な考慮はいりませんが、いくつかの注意事項があり、これについては tf.function ガイドのほか、Autograph 完全リファレンスが役立ちます。

高速化の確認

tensor-using 関数を tf.function でラッピングするだけでは、コードは高速化しません。単一のマシンで数回呼び出された小さな関数では、グラフまたはグラフの一部の呼び出しにかかるオーバーヘッドによってランタイムが占有されてしまうことがあります。また、GPU 大きな負荷をかける畳み込みのスタックなど、計算のほとんどがすでにアクセラレータで発生している場合は、グラフの高速化をあまり確認できません。

複雑な計算については、グラフによって大幅な高速化を得ることができます。これは、グラフが Python からデバイスへの通信や一部の高速化の実装を減らすためです。

次のコードは、小さな密のレイヤーでの数回の実行にかかる時間を計測します。

# Create an oveerride model to classify pictures
class SequentialModel(tf.keras.Model):
  def __init__(self, **kwargs):
    super(SequentialModel, self).__init__(**kwargs)
    self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
    self.dense_1 = tf.keras.layers.Dense(128, activation="relu")
    self.dropout = tf.keras.layers.Dropout(0.2)
    self.dense_2 = tf.keras.layers.Dense(10)

  def call(self, x):
    x = self.flatten(x)
    x = self.dense_1(x)
    x = self.dropout(x)
    x = self.dense_2(x)
    return x

input_data = tf.random.uniform([60, 28, 28])

eager_model = SequentialModel()
graph_model = tf.function(eager_model)

print("Eager time:", timeit.timeit(lambda: eager_model(input_data), number=10000))
print("Graph time:", timeit.timeit(lambda: graph_model(input_data), number=10000))
Eager time: 7.648914662999687
Graph time: 2.6870547049998095

多層型関数

関数をトレースする場合、多層型Function オブジェクトを作成します。多層型関数は Pythonコーラブルで、1つの API の背後にあるいくつかの具象関数グラフをカプセル化します。

この Function は、あらゆる dtypes と形状に使用できます。新しい引数シグネチャでそれを呼び出すたびに、元の関数が新しい引数で再トレースされます。Function は、そのトレースに対応する tf.Graphconcrete_function に格納します。関数がすでにそのような引数でトレースされている場合は、トレース済みのグラフが取得されます。

概念的に、次のようになります。

  • tf.Graph は計算を説明する未加工のポータブルなデータ構造である
  • Function は、ConcreteFunctions のキャッシュ、トレース、およびディスパッチャーである
  • ConcreteFunction は、Python からグラフを実行できるグラフの Eager 対応ラッパーである

多層型関数の検査

a_function を検査できます。これはPython 関数 my_function に対して tf.function を呼び出した結果です。この例では、3 つの引数で a_function を呼び出すことで、3 つの具象関数を得られています。

print(a_function)

print("Calling a `Function`:")
print("Int:", a_function(tf.constant(2)))
print("Float:", a_function(tf.constant(2.0)))
print("Rank-1 tensor of floats", a_function(tf.constant([2.0, 2.0, 2.0])))
<tensorflow.python.eager.def_function.Function object at 0x7f4d79096908>
Calling a `Function`:
Int: tf.Tensor(1, shape=(), dtype=int32)
Float: tf.Tensor(1.0, shape=(), dtype=float32)
Rank-1 tensor of floats tf.Tensor([1. 1. 1.], shape=(3,), dtype=float32)

# Get the concrete function that works on floats
print("Inspecting concrete functions")
print("Concrete function for float:")
print(a_function.get_concrete_function(tf.TensorSpec(shape=[], dtype=tf.float32)))
print("Concrete function for tensor of floats:")
print(a_function.get_concrete_function(tf.constant([2.0, 2.0, 2.0])))
Inspecting concrete functions
Concrete function for float:
ConcreteFunction my_function(x)
  Args:
    x: float32 Tensor, shape=()
  Returns:
    float32 Tensor, shape=()
Concrete function for tensor of floats:
ConcreteFunction my_function(x)
  Args:
    x: float32 Tensor, shape=(3,)
  Returns:
    float32 Tensor, shape=(3,)

# Concrete functions are callable
# Note: You won't normally do this, but instead just call the containing `Function`
cf = a_function.get_concrete_function(tf.constant(2))
print("Directly calling a concrete function:", cf(tf.constant(2)))
Directly calling a concrete function: tf.Tensor(1, shape=(), dtype=int32)

この例では、スタックの非常に奥を調べています。具体的にトレースを管理していない限り、通常は、ここに示されるように具象関数を呼び出す必要はありません。

Eager execution でのデバッグ

スタックトレースが長い場合、特に tf.Graph または with tf.Graph().as_default() の参照が含まれる場合、グラフコンテキストで実行している可能性があります。TensorFlow のコア関数は Keras の model.fit() などのグラフコンテキストを使用します。

Eager execution をデバッグする方がはるかに簡単であることがよくあります。スタックトレースは比較的に短く、理解しやすいからです。

グラフのデバックが困難な場合は、Eager execution に戻ってデバックすることができます。

Eager で実行していることを確認するには、次を行います。

  • メソッドとレイヤーを直接コーラブルとして呼び出す

  • Keras compile/fit を使用している場合、コンパイル時に model.compile(run_eagerly=True) を使用する

  • tf.config.experimental_run_functions_eagerly(True) でグローバル実行モードを設定する

run_eagerly=True を使用する

# Define an identity layer with an eager side effect
class EagerLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(EagerLayer, self).__init__(**kwargs)
    # Do some kind of initialization here

  def call(self, inputs):
    print("\nCurrently running eagerly", str(datetime.now()))
    return inputs
# Create an override model to classify pictures, adding the custom layer
class SequentialModel(tf.keras.Model):
  def __init__(self):
    super(SequentialModel, self).__init__()
    self.flatten = tf.keras.layers.Flatten(input_shape=(28, 28))
    self.dense_1 = tf.keras.layers.Dense(128, activation="relu")
    self.dropout = tf.keras.layers.Dropout(0.2)
    self.dense_2 = tf.keras.layers.Dense(10)
    self.eager = EagerLayer()

  def call(self, x):
    x = self.flatten(x)
    x = self.dense_1(x)
    x = self.dropout(x)
    x = self.dense_2(x)
    return self.eager(x)

# Create an instance of this model
model = SequentialModel()

# Generate some nonsense pictures and labels
input_data = tf.random.uniform([60, 28, 28])
labels = tf.random.uniform([60])

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

まず、Eager を使用せずにモデルをコンパイルします。モデルはトレースされません。名前にも関わらず、compile は、損失関数、最適化、およびトレーニングパラメータのセットアップしか行いません。

model.compile(run_eagerly=False, loss=loss_fn)

ここで、fit を呼び出し、関数がトレース(2 回)されると Eager 効果が実行しなくなるのを確認します。

model.fit(input_data, labels, epochs=3)
Epoch 1/3

Currently running eagerly 2021-02-12 21:38:32.736480

Currently running eagerly 2021-02-12 21:38:32.857779
2/2 [==============================] - 0s 3ms/step - loss: 1.9410
Epoch 2/3
2/2 [==============================] - 0s 2ms/step - loss: 0.0036
Epoch 3/3
2/2 [==============================] - 0s 2ms/step - loss: 0.0013

<tensorflow.python.keras.callbacks.History at 0x7f4cc00ec668>

ただし、エポックを 1 つでも Eager で実行すると、Eager の副次的作用が 2 回現れます。

print("Running eagerly")
# When compiling the model, set it to run eagerly
model.compile(run_eagerly=True, loss=loss_fn)

model.fit(input_data, labels, epochs=1)
Running eagerly

Currently running eagerly 2021-02-12 21:38:33.068015
1/2 [==============>...............] - ETA: 0s - loss: 6.9482e-04
Currently running eagerly 2021-02-12 21:38:33.089168
2/2 [==============================] - 0s 12ms/step - loss: 4.7885e-04

<tensorflow.python.keras.callbacks.History at 0x7f4cc03748d0>

experimental_run_functions_eagerly を使用する

また、すべてを Eager で実行するよにグローバルに設定することができます。これは、トレースし直した場合にのみ機能することに注意してください。トレースされた関数は、トレースされたままとなり、グラフとして実行します。

# Now, globally set everything to run eagerly
tf.config.experimental_run_functions_eagerly(True)
print("Run all functions eagerly.")

# First, trace the model, triggering the side effect
polymorphic_function = tf.function(model)

# It was traced...
print(polymorphic_function.get_concrete_function(input_data))

# But when you run the function again, the side effect happens (both times).
result = polymorphic_function(input_data)
result = polymorphic_function(input_data)
WARNING:tensorflow:From <ipython-input-1-6bee22f96a72>:2: experimental_run_functions_eagerly (from tensorflow.python.eager.def_function) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.run_functions_eagerly` instead of the experimental version.
Run all functions eagerly.

Currently running eagerly 2021-02-12 21:38:33.117223
ConcreteFunction function(self)
  Args:
    self: float32 Tensor, shape=(60, 28, 28)
  Returns:
    float32 Tensor, shape=(60, 10)

Currently running eagerly 2021-02-12 21:38:33.121492

Currently running eagerly 2021-02-12 21:38:33.123248

# Don't forget to set it back when you are done
tf.config.experimental_run_functions_eagerly(False)

トレースとパフォーマンス

トレースにはある程度のオーバーヘッドがかかります。小さな関数のトレースは素早く行えますが、大規模なモデルであればかなりの時間がかかる場合があります。パフォーマンスが上昇するとこの部分の時間は迅速に取り戻されますが、大規模なモデルのトレーニングの最初の数エポックでは、トレースによって遅延が発生する可能性があることに注意しておくことが重要です。

モデルの規模に関係なく、頻繁にトレースするのは避けたほうがよいでしょう。tf.function ガイドのこのセクションでは、入力仕様を設定し、テンソル引数を使用して再トレースを回避する方法について説明しています。フォーマンスが異常に低下している場合は、誤って再トレースしていないかどうかを確認することをお勧めします。

eager-only の副次的効果(Python 引数の出力など)を追加して、関数がいつトレースされているかを確認できます。ここでは、新しい Python 引数が常に再トレースをトリガするため、余分な再トレースが発生していることを確認できます。

# Use @tf.function decorator
@tf.function
def a_function_with_python_side_effect(x):
  print("Tracing!")  # This eager
  return x * x + tf.constant(2)

# This is traced the first time
print(a_function_with_python_side_effect(tf.constant(2)))
# The second time through, you won't see the side effect
print(a_function_with_python_side_effect(tf.constant(3)))

# This retraces each time the Python argument chances
# as a Python argument could be an epoch count or other
# hyperparameter
print(a_function_with_python_side_effect(2))
print(a_function_with_python_side_effect(3))
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
tf.Tensor(11, shape=(), dtype=int32)
Tracing!
tf.Tensor(6, shape=(), dtype=int32)
Tracing!
tf.Tensor(11, shape=(), dtype=int32)

次のステップ

より詳しい説明については、tf.function API リファレンスページとガイドを参照してください。