TensorFlow 2.4 : ガイド : 基本 – グラフと tf.functions へのイントロダクション (翻訳/解説)

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

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

• Introduction to graphs and tf.functions

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

 

 

ガイド : 基本 – グラフと tf.functions へのイントロダクション

概要

このガイドは TensorFlow がどのように動作するかを見るために TensorFlow と Keras の表面下を進みます。代わりに直ちに Keras から始めることを望む場合には、Keras ガイドのコレクション を見てください。

このガイドでは、グラフを得るために TensorFlow がどのようにコードに単純な変更を行なうことを貴方に可能にするか、そしてそれらがどのようにストアされて表されるか、そしてモデルをアクセラレートしてエキスポートするためにそれらをどのように利用できるかのコアを見るでしょう。

これは短い形式のイントロダクションです ; これらのコンセプトへの完全なイントロダクションについては、tf.function ガイド を見てください。

 

グラフとは何でしょうか？

前の 3 つのガイドでは、TensorFlow が eagerly に動作することを見ました。これは TensorFlow 演算は Python により実行され、演算毎に、そして結果を Python に返し戻します。eager TensorFlow は GPU を利用し、変数、tensor そして演算さえも GPU と TPU 上に配置することを可能にします。デバッグすることもまた容易です。

あるユーザにとっては、決して Python を離れる必要がなく望むこともないかもしれません。

けれども、Python で TensorFlow を演算毎に実行することは (さもなければ利用可能な) 多くのアクセラレーションを妨げます。Python から tensor 計算を抽出できれば、それらをグラフにできます。

グラフは計算のユニットを表す tf.Operation オブジェクト; そして演算間をフローするデータのユニットを表す tf.Tensor オブジェクトのセットを含むデータ構造です。それらは tf.Graph コンテキストで定義されます。これらのグラフはデータ構造ですから、それらは元の Python コードなしに総てセーブされ、実行されてリストアされます。

これは単純な 2-層グラフが TensorBoard で可視化されたときどのように見えるかです。

 

グラフの恩恵

グラフにより、非常に大きな柔軟性を持ちます。モバイルアプリケーション、埋め込みデバイス、そしてバックエンドサーバのような Python インタープリタを持たない環境で TensorFlow グラフを利用できます。TensorFlow はグラフを Python からエクスポートするときセーブされたモデルのための形式として使用します。

グラフはまた容易に最適化されて、コンパイラに以下のような変換を行なうことを可能にします :

• 計算で定数ノードを折り畳む (“constant folding”) ことにより tensor の値を静的に推論します。
• 独立な計算のサブパートを分離してそれらをスレッドかデバイス間で分割します。
• 共通部分式の除去により算術演算を単純化します。

これと他のスピードアップを遂行するために最適化全体のシステム, Grappler があります。

手短に言えば、グラフは極めて有用でそして TensorFlow を マルチデバイス 上で 高速 に、並列 に、そして効率的に実行させます。

けれども、便利のために Python で機械学習モデル (or 他の計算) を定義して、それからグラフを必要とするときに自動的に構築することを貴方は依然として望むでしょう。

 

グラフをトレースする

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-ized 関数はそれらの Python 等値と同じ動作をする Python callable です。それらは特定のクラス (python.eager.def_function.Function) を持ちますが、貴方に対してはそれらは単に非-traced バージョンとして動作します。

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 は standardizing loop constructs, unrolling と 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.]

Python がどのように TensorFlow ops にどのように変換されるかを見るために Autograph 変換を直接呼び出すことができます。これは、殆ど、可読ではありませんが、変換を見ることができます。

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

Autograh は if-then 節, loops, break, return, continue 等々を自動的に変換します。

大抵の場合、Autograph は特別な考慮なく動作するでしょう。けれども幾つかの注意事項があり、そして 完全な autograph リファレンス に加えて tf.function ガイド がここでは手助けできます。

 

スピードアップを見る

tf.function で tensor を使用する関数を単にラッピングしても貴方のコードを自動的にはスピードアップしません。単一のマシン上何回か呼び出される関数のため、グラフかグラフ断片を呼び出すオーバーヘッドが実行時間を支配するかもしれません。また、GPU heavy な畳込みのスタックのように、計算の殆どが既にアクセラレータ上で起きている場合、グラフのスピードアップは大きくないでしょう。

複雑な計算のために、グラフは本質的なスピードアップを提供できます。これはグラフが Python-to-デバイス通信を削減して何某かのスピードアップを遂行するからです。

下のサンプル内のように、多くの小さい層を実行するときスピードアップは最も明白です。

# 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))
    # Add a lot of small layers
    num_layers = 100
    self.my_layers = [tf.keras.layers.Dense(64, activation="relu")
                      for n in range(num_layers)]
    self.dropout = tf.keras.layers.Dropout(0.2)
    self.dense_2 = tf.keras.layers.Dense(10)

  def call(self, x):
    x = self.flatten(x)
    for layer in self.my_layers:
      x = layer(x)
    x = self.dropout(x)
    x = self.dense_2(x)
    return x

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

eager_model = SequentialModel()

# Don't count the time for the initial build.
eager_model(input_data)
print("Eager time:", timeit.timeit(lambda: eager_model(input_data), number=100))

Eager time: 1.9846877160000531

(訳注: 下は訳者による実験)

Eager time: 2.2071227299999805

# Wrap the call method in a `tf.function`
graph_model = SequentialModel()
graph_model.call = tf.function(graph_model.call)

# Don't count the time for the initial build and trace.
graph_model(input_data)
print("Graph time:", timeit.timeit(lambda: graph_model(input_data), number=100))

Graph time: 0.32060739400003513

(訳注: 下は訳者による実験)

Graph time: 0.3933808909999925

 

Polymorphic 関数

関数をトレースするとき、poymorphic (多態な) Function オブジェクトを作成します。polymorphic 関数は一つの API の裏で幾つかの具体的なグラフをカプセル化する Python callable です。

この Function を総ての異なる種類の dtype と shape 上で使用できます。それを新しい引数シグネチャで呼び出すたびに、元の関数は新しい引数で再トレースされます。それから Function は concrete_function でそのトレースに対応する tf.Graph をストアします。関数が既にその種類の引数でトレースされあている場合、事前トレースされたグラフを単に得るだけです。

概念的には、それから :

• tf.Graph は計算を記述する raw で可搬なデータ構造です。
• Function は ConcreteFunction に渡るキャッシュ、トレーシング、ディスパッチャーです。
• ConcreteFunction はグラフ回りの eager-互換ラッパーで、これは Python からグラフを貴方に実行させます。

 

polymorphic 関数を調査する

a_function を調査することができます、これは Python 関数 my_function 上で tf.function を呼び出した結果です。このサンプルでは、a_function を 3 つの種類の引数で呼び出し 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 0x7f124ab79588>
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 実行に戻る

貴方自身が長いスタックトレースを見ていることを見出すかもしれません、特に tf.Graph を参照するか tf.Graph().as_default() によるものは。これはグラフコンテキストで実行している可能性が高いことを意味します。Keras の model.fit() のような、TensorFlow のコア関数はグラフコンテキストを使用します。

eager 実行をデバッグすることはしばしば遥かに容易です。スタックトレースは比較的短くて理解するのが容易なはずです。

グラフがデバッグをトリッキーにする状況では、デバッグするために eager 実行の使用に戻ることができます。

ここに貴方が eagerly に実行していることを確かにできる方法があります :

• モデルと層を callable として直接呼び出す。
• Keras compile/fit を使用するとき、コンパイル時に model.compile(run_eagerly=True) を使用します。
• tf.config.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 を呼び出して関数がトレース (二度) されてから eager エフェクトが決して再度実行されないことを見ます。

model.fit(input_data, labels, epochs=3)

Epoch 1/3

Currently running eagerly 2020-12-15 02:32:22.321423

Currently running eagerly 2020-12-15 02:32:22.435796
2/2 [==============================] - 0s 3ms/step - loss: 1.6631
Epoch 2/3
2/2 [==============================] - 0s 2ms/step - loss: 0.0011
Epoch 3/3
2/2 [==============================] - 0s 3ms/step - loss: 5.7225e-04

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

(訳注: 下は訳者による実験)

Epoch 1/3

Currently running eagerly 2020-12-25 05:37:06.506613

Currently running eagerly 2020-12-25 05:37:06.633045
2/2 [==============================] - 0s 5ms/step - loss: 1.5245
Epoch 2/3
2/2 [==============================] - 0s 4ms/step - loss: 0.0022
Epoch 3/3
2/2 [==============================] - 0s 3ms/step - loss: 0.0015
<tensorflow.python.keras.callbacks.History at 0x7f37a57584e0>

eager で単一エポックを実行する場合でさえも、けれども、eager サイド・エフェクトを二度見ることができます。

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 2020-12-15 02:32:22.640829
1/2 [==============>...............] - ETA: 0s - loss: 4.5348e-04
Currently running eagerly 2020-12-15 02:32:22.660222
2/2 [==============================] - 0s 13ms/step - loss: 3.1242e-04

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

(訳注: 下は訳者による実験)

Running eagerly

Currently running eagerly 2020-12-25 05:43:36.517554
1/2 [==============>...............] - ETA: 0s - loss: 0.0020
Currently running eagerly 2020-12-25 05:43:36.544462
2/2 [==============================] - 0s 14ms/step - loss: 0.0014
<tensorflow.python.keras.callbacks.History at 0x7f37a561b2b0>

 

run_functions_eagerly を使用する

eagerly に実行するためグローバルに総てを設定することもできます。これは polymorphic 関数のトレースされた関数を迂回して元の関数を直接呼び出すスイッチです。これをデバッグのために利用できます。

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

# Create a polymorphic function
polymorphic_function = tf.function(model)

print("Tracing")
# This does, in fact, trace the function
print(polymorphic_function.get_concrete_function(input_data))

print("\nCalling twice eagerly")
# When you run the function again, you will see the side effect
# twice, as the function is running eagerly.
result = polymorphic_function(input_data)
result = polymorphic_function(input_data)

Run all functions eagerly.
Tracing

Currently running eagerly 2020-12-25 09:29:29.001118
ConcreteFunction function(self)
  Args:
    self: float32 Tensor, shape=(60, 28, 28)
  Returns:
    float32 Tensor, shape=(60, 10)

Calling twice eagerly

Currently running eagerly 2020-12-25 09:29:29.007092

Currently running eagerly 2020-12-25 09:29:29.009134

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

WARNING:tensorflow:From <ipython-input-21-782fe9ce7b18>: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.

 

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

トレーシングは何某かのオーバーヘッドがかかります。小さい関数のトレースは迅速ですが、大規模なモデルはトレースするために顕著な実測時間がかかる可能性があります。この投資は通常はパフォーマンス・ブーストにより素早く払い戻されますが、任意の大規模モデル訓練の最初の数エポックはトレーシングのために遅くなる可能性があることを知っていることは重要です。

貴方のモデルがどれくらい大規模であろうと、頻繁にトレースすることを回避することを望みます。tf.function ガイドのこのセクション は再トレーシングを回避するためにどのように入力仕様を設定して tensor 引数を使用するかを議論します。異常に貧弱なパフォーマンスを得ていることを見出す場合、再トレースを偶々していないかを確認するのが良いです。

(Python 引数をプリントするような) eager-only サイド・エフェクトを追加できますので、関数がいつトレースされるかを見ることができます。ここでは、余分な (= extra) 再トレーシングを見ます、何故ならば新しい 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 changes,
# 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)

以上