tf.functionの再トレースによる訓練の低速化について確かめる

---

TensorFlow2.0において、tf.functionを使うと計算が高速化することはよく知られていますが、その代償として入力のshapeが可変の場合に「再トレース」が発生し訓練が低速化することがあります。再トレースが頻発するケースに遭遇したので、どの程度悪化するものなのか確かめてみました。

TensorFlowはもともとDefine-and-run

これを知るにはTensorFlowの歴史的な背景を知っておくと理解がしやすいかもしれません。TensorFlowは特にVer1系で「Define-and-run」という計算グラフのスタイルを取っていました。一方でChainerやPyTorchは「Define-by-run」というスタイルを取っています。

Define-and-runとDefine-by-runの違いは何かというと、TensorFlow1系のようなDefine-and-runは計算の流れ（計算グラフ）をガッチリ固めてから実行するという方式でした。これはモデルを柔軟性が落ちる一方で速度を出しやすいというメリットがあります。ディープラーニングのフレームワークはGPUの低レベルの計算をあくまでラップして使いやすいようにしてるにすぎないので、低レベルのコマンドを固定できるというのは単純にパフォーマンス面では優位に立ちやすいでしょう。

一方のDefine-by-runは、最初にガチッと固めずに（forward関数を実行させながら）計算するものです。こちらはforward関数の中にif文を入れたり、コーディングやモデルの柔軟性を確保しやすいというメリットがあります。一方で、動的なモデルであることは変わらないので、本質的に速度が犠牲になりがちな側面はあります。ただし、最近のフレームワークはとても優秀で相当工夫が入っているので、ちょっとやそっとでは速度は落ちません。

一方で、TF2.0になってからEager Executionという動的実行可能なシステムが本格的に導入されました。これは、Define-by-runとDefine-and-runの速度の差がほとんどなくなってきて、より柔軟性が欲しいという需要が高まってきたためです（PyTorchが一気に伸びた理由の1つでもあります）。つまり、もともとTensorFlowはDefine-and-runであったけれども、Eager ExecutionによってDefine-by-run的なこともできるようになったというのが現状です。

Eager Executionの中でのtf.function

公式ドキュメントではこちらに詳しく書かれている内容です（日本語）

tf.function で性能アップ

Eager Executionが前提となったTF2.0でのtf.functionとはどういう扱いかというと、tf.functionの中で囲まれた部分を1つの固定のグラフと見る、ざっくりとした言い方では、「tf.functionで囲まれた部分をDefine-and-runのように見る」というイメージです。

ディープラーニングのフレームワークとしてのDefine-and-runとDefine-by-runの速度差はほとんどなくても、固定のグラフに対してtf.functionの有無による速度差は非常に大きく現れます（特に微分計算が入る場合）。

入力サイズが固定の場合のtf.functionのパワー

例えば、VGG16の224×224サイズの入力に対し学習を行ってみましょう。ベンチマークを取るだけなのでただの乱数で学習させています。

環境：TF2.1.0、RTX 2080 Ti、CUDA10.1、Windows10

import tensorflow as tf
import tensorflow.keras.layers as layers
import numpy as np
from tqdm import tqdm

def static_vgg():
    model = tf.keras.applications.VGG16(include_top=False, input_shape=(224, 224, 3))
    optim = tf.keras.optimizers.Adam()

    @tf.function
    def on_batch(inputs, y_true):
        with tf.GradientTape() as tape:
            y_pred = model(inputs, training=True)
            loss = tf.reduce_mean((y_true - y_pred)** 2, axis=(1, 2, 3))
        grads = tape.gradient(loss, model.trainable_weights)
        optim.apply_gradients(zip(grads, model.trainable_weights))
        return loss

    for i in tqdm(range(1000)):
        x = np.random.randn(16, 224, 224, 3).astype(np.float32)
        y_true = np.random.randn(16, 7, 7, 512).astype(np.float32)        
        on_batch(x, y_true)

if __name__ == "__main__":
    static_vgg()

実行時間をon_batchのtf.functionの有無で比較します。結果は、

• tf.functionあり：1分52秒
• tf.functionなし：3分27秒

tf.functionの有無で2倍近い差になりました。これは無視できない差です。

なぜtf.functionを入れると高速化するのかというと、tf.functionを入れると最初に定義した計算グラフを使い回すことができるからです（この場合は入力サイズが固定なので簡単に使い回せる）。ここでは詳しくは述べませんが、tf.functionを入れない場合、モデル定義で（call関数など）中間層の値を見ることができるので、デバッグ用にtf.functionを消して本番用にtf.functionを入れるというような使い分けもできます。

入力サイズが可変の場合の問題

問題は入力サイズが可変の場合です。このケースではtf.functionを入れればいいかというとそう単純ではありません。なぜなら、同じVGG16でも入力サイズが256×256と64×64で出てくるテンソルのshapeが異なるからです。

入力サイズを可変にするための準備として、VGGのinput_shapeを「(224, 224, 3)」のような固定値から「(None, None, 3)」のようなNoneを使ったshapeに変更します。Noneを次元数の指定に使うことはKeras APIではそう珍しいことではなくて、バッチサイズの軸はデフォルトでNoneです。なので、Noneの次元が増えただけということです。やり方は簡単で、

model = tf.keras.applications.VGG16(include_top=False, input_shape=(None, None, 3))

とすればいいだけです。このモデルのsummaryは、

（略）
_________________________________________________________________
block5_conv1 (Conv2D)        (None, None, None, 512)   2359808
_________________________________________________________________
block5_conv2 (Conv2D)        (None, None, None, 512)   2359808
_________________________________________________________________
block5_conv3 (Conv2D)        (None, None, None, 512)   2359808
_________________________________________________________________
block5_pool (MaxPooling2D)   (None, None, None, 512)   0
=================================================================
Total params: 14,714,688
Trainable params: 14,714,688
Non-trainable params: 0

と今まで固定値で表示されていたのがNoneで出るようになるだけです（モデル内にReshape操作が入ると扱いが面倒になるんですけどね）。

以下のコードは縦横解像度が64～256の乱数であるケースで、ミニバッチごとに解像度が変わります。

def dynamic_vgg_notrain():
    model = tf.keras.applications.VGG16(include_top=False, input_shape=(None, None, 3))

    #@tf.function(input_signature=(tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.float32),))
    @tf.function
    def on_batch(inputs):
        return model(inputs)

    for i in tqdm(range(10)):
        h = np.random.randint(64, 256)
        w = np.random.randint(64, 256)
        x = np.random.randn(16, h, w, 3).astype(np.float32)
        y = on_batch(x)

if __name__ == "__main__":
    dynamic_vgg_notrain()

しかし、先程のように単にtf.funcitonを入れて実行すると、次のようなメッセージで怒られます。

10 out of the last 10 calls to <function dynamic_vgg_notrain.<locals>.on_batch at 0x000002CD87D78E18> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.

ここでこの記事のタイトルにもなっている再トレース（Retracing）の問題が出てきます。エラー文の内容は「再トレースはコストが高いから以下の記事読んで対策しろ」という趣旨ですね。入力サイズが可変になることと再トレースがおこることは関係していて、公式ドキュメントによると、

他方で、TensorFlow の計算グラフでは、dtype と shape の次元が静的であることが必要です。tf.function は、正しい計算グラフを生成するために必要なときには関数を再トレースして、このギャップをつなぐ役割を果たします。

これは裏を返せば、入力サイズが可変だとshapeが動的になるので、再トレースが発生するということになります。ただし対策方法はあります。

対策方法その1：tf.functionを入れない

後で確かめますが、意外とこれがわかりやすく、入力サイズが可変のケースでは下手に書いたtf.functionよりも高速に動くことがあります。tf.functionを入れないことで、モデルの柔軟性が上がり、訓練中にNumpy操作でごちゃごちゃ書いてても通るようになります（これが望ましいかどうかは別として）

対策方法その2 tf.functionの引数input_signatureで入力shapeと型を明示する

公式ドキュメントにある以下のテクニックを使います

トレースの動作を制御するためには、下記のようなテクニックを使います。
… 計算グラフの呼び出し時に1回だけトレースを行うには、 input_signature を指定して tf.function を呼び出す。

ここでポイントなのはinput_signatureのshapeはNoneでもOKなので、

    @tf.function(input_signature=(tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.float32),))
    def on_batch(inputs):
        return model(inputs)

という書き方をすれば再トレースを発生させないようにすることが可能です。input_signatureを指定して実行すると、再トレースに関する警告メッセージがずらーっと表示されることはなくなります。

入力サイズが可変かつ訓練なしの場合の実験

tf.functionの効果は微分があるときに大きく現れますが、可変の場合は訓練なし（Fore-propだけ）でも目に見える差が出ます。

def dynamic_vgg_notrain():
    model = tf.keras.applications.VGG16(include_top=False, input_shape=(None, None, 3))

    @tf.function(input_signature=(tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.float32),))
    def on_batch(inputs):
        return model(inputs)

    for i in tqdm(range(1000)):
        h = np.random.randint(64, 256)
        w = np.random.randint(64, 256)
        x = np.random.randn(16, h, w, 3).astype(np.float32)
        y = on_batch(x)

if __name__ == "__main__":
    dynamic_vgg_notrain()

このon_batchのデコレーターを次の3種類で実験してみます。

1. デコレーター（tf.function）なし
2. tf.functionのデコレーターをつけただけ、引数はなし
3. tf.functionのデコレーターを、input_signatureの引数を指定してつける

3回計測した結果は次のようになりました

1. tf.functionなし：7分58秒、8分09秒、8分04秒
2. tf.functionあり、input_signatureなし：8分46秒、8分55秒、8分37秒
3. tf.functionあり、input_signatureあり：7分43秒、7分49秒、7分52秒

Fore-propだけでも、tf.functionありinput_signatureありが最速になりました。tf.functionなしでも若干遅いぐらいでとどまりそうです。しかし、再トレースが多発するinput_signatureなしのtf.functionでは明らかに遅くなっています。これが再トレースによる低速化です。

入力サイズが可変かつ訓練ありの場合の実験

訓練あり（Back-prop込み）の場合はどうでしょうか？以下のコードを同様に3パターンで実験してみます。時間短縮のために繰り返し回数を1000から300に減らしています。

def dynamic_vgg_withtrain():
    model = tf.keras.applications.VGG16(include_top=False, input_shape=(None, None, 3))
    optim = tf.keras.optimizers.Adam()

    @tf.function(input_signature=(tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.float32),
                                  tf.TensorSpec(shape=[None, None, None, 512], dtype=tf.float32)))
    def on_batch(inputs, y_true):
        with tf.GradientTape() as tape:
            y_pred = model(inputs, training=True)
            loss = tf.reduce_mean((y_true - y_pred)** 2, axis=(1, 2, 3))
        grads = tape.gradient(loss, model.trainable_weights)
        optim.apply_gradients(zip(grads, model.trainable_weights))
        return loss

    for i in tqdm(range(300)):
        h = np.random.randint(64, 256)
        w = np.random.randint(64, 256)
        x = np.random.randn(16, h, w, 3).astype(np.float32)
        y_true = np.random.randn(16, h // 32, w // 32, 512).astype(np.float32)        
        on_batch(x, y_true)

結果は、

1. tf.functionなし：7分12秒
2. tf.functionあり、input_signatureなし：7分41秒
3. tf.functionあり、input_signatureあり：6分57秒

となりました。Back-propありでもVGG16では、tf.functionありinput_signatureありが最速になりました。

もっと深いネットワークの場合（ResNet152）

より深いネットワークになると再トレースのコストがもっと深刻になるのではないか？と思ったため、VGG16よりレイヤー数が遥かに多いResNet152で確かめてみました。

def dynamic_resnet():
    model = tf.keras.applications.ResNet152V2(include_top=False, input_shape=(None, None, 3))
    optim = tf.keras.optimizers.Adam()

    @tf.function(input_signature=(tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.float32),
                                  tf.TensorSpec(shape=[None, None, None, 2048], dtype=tf.float32)))
    def on_batch(inputs, y_true):
        with tf.GradientTape() as tape:
            y_pred = model(inputs, training=True)
            loss = tf.reduce_mean((y_true - y_pred)** 2, axis=(1, 2, 3))
        grads = tape.gradient(loss, model.trainable_weights)
        optim.apply_gradients(zip(grads, model.trainable_weights))
        return loss

    for i in tqdm(range(300)):
        h = np.random.randint(64, 256)
        w = np.random.randint(64, 256)
        x = np.random.randn(16, h, w, 3).astype(np.float32)
        y_true = np.random.randn(16, int(np.ceil(h / 32)), int(np.ceil(w / 32)), 2048).astype(np.float32)
        on_batch(x, y_true)

ResNet152の場合はかなり結果が変わり、

1. tf.functionなし：3分36秒
2. tf.functionあり、input_signatureなし：36分10秒
3. tf.functionあり、input_signatureあり：4分37秒

2が再トレース多発するケースですが、再トレースのコストが非常に深刻になっているのがわかります。3～4分で収まる訓練が36分かかってしまいます。

ResNet152の場合は、input_signatureありのtf.functionよりも、tf.functionそのものを切ったほうが速いという結果になりました。どちらが速いかはケースバイケースでしょうか。コーディングの柔軟性を考えると、入力サイズが可変のケースでは、一部のネットワークでは多少遅くなってもtf.functionを切ったほうが良いということもあり得るでしょう。

まとめ

これまでの計測結果をまとめます。

ネットワーク	訓練	入力サイズ	繰り返し回数	tf.functionなし	tf.functionあり、input_signatureなし	tf.functionあり、input_signatureあり
VGG16	○	224固定	1000	3:27	1:52	–
VGG16	☓	64-256可変	1000	8:03	8:46	7:48
VGG16	○	64-256可変	300	7:12	7:41	6:57
ResNet152	○	64-256可変	300	3:36	36:10	4:37

入力サイズが可変のケースは遅いのは仕方ないですね。固定よりも10倍前後遅くなっています。

ただ、可変のケースは可変のケースで便利な点もあるので（入力画像のアスペクト比の問題とか）、可変で使うにしても再トレースは起こすメリットがまったくないということは言えるでしょう。

tf.functionなしにするか、tf.function入れて再トレース起こさないようにinput_signatureを指定するかはモデルによって異なる結果となりました。

---

---

Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内

---