こしあん
2020-05-18

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

Pocket
LINEで送る
Delicious にシェア

1.9k{icon} {views}



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の中の人が運営しているサークル「じゅ~しぃ~すくりぷと」の本のご案内

技術書コーナー

【新刊】インフィニティNumPy――配列の初期化から、ゲームの戦闘、静止画や動画作成までの221問

「本当の実装力を身につける」ための221本ノック――
機械学習(ML)で避けて通れない数値計算ライブラリ・NumPyを、自在に活用できるようになろう。「できる」ための体系的な理解を目指します。基礎から丁寧に解説し、ディープラーニング(DL)の難しいモデルで遭遇する、NumPyの黒魔術もカバー。初心者から経験者・上級者まで楽しめる一冊です。問題を解き終わったとき、MLやDLなどの発展分野にスムーズに入っていけるでしょう。

本書の大きな特徴として、Pythonの本でありがちな「NumPyとML・DLの結合を外した」点があります。NumPyを理解するのに、MLまで理解するのは負担が大きいです。本書ではあえてこれらの内容を書いていません。行列やテンソルの理解に役立つ「従来の画像処理」をNumPyベースで深く解説・実装していきます。

しかし、問題の多くは、DLの実装で頻出の関数・処理を重点的に取り上げています。経験者なら思わず「あー」となるでしょう。関数丸暗記では自分で実装できません。「覚える関数は最小限、できる内容は無限大」の世界をぜひ体験してみてください。画像編集ソフトの処理をNumPyベースで実装する楽しさがわかるでしょう。※紙の本は電子版の特典つき

モザイク除去から学ぶ 最先端のディープラーニング

「誰もが夢見るモザイク除去」を起点として、機械学習・ディープラーニングの基本をはじめ、GAN(敵対的生成ネットワーク)の基本や発展型、ICCV, CVPR, ECCVといった国際学会の最新論文をカバーしていく本です。
ディープラーニングの研究は発展が目覚ましく、特にGANの発展型は市販の本でほとんどカバーされていない内容です。英語の原著論文を著者がコードに落とし込み、実装を踏まえながら丁寧に解説していきます。
また、本コードは全てTensorFlow2.0(Keras)に対応し、Googleの開発した新しい機械学習向け計算デバイス・TPU(Tensor Processing Unit)をフル活用しています。Google Colaboratoryを用いた環境構築不要の演習問題もあるため、読者自ら手を動かしながら理解を深めていくことができます。

AI、機械学習、ディープラーニングの最新事情、奥深いGANの世界を知りたい方にとってぜひ手にとっていただきたい一冊となっています。持ち運びに便利な電子書籍のDLコードが付属しています。

「おもしろ同人誌バザールオンライン」で紹介されました!(14:03~) https://youtu.be/gaXkTj7T79Y?t=843

まとめURL:https://github.com/koshian2/MosaicDeeplearningBook
A4 全195ページ、カラー12ページ / 2020年3月発行

Shikoan's ML Blog -Vol.1/2-

累計100万PV超の人気ブログが待望の電子化! このブログが電子書籍になって読みやすくなりました!

・1章完結のオムニバス形式
・機械学習の基本からマニアックなネタまで
・どこから読んでもOK
・何巻から読んでもOK

・短いものは2ページ、長いものは20ページ超のものも…
・通勤・通学の短い時間でもすぐ読める!
・読むのに便利な「しおり」機能つき

・全巻はA5サイズでたっぷりの「200ページオーバー」
・1冊にたっぷり30本収録。1本あたり18.3円の圧倒的コストパフォーマンス!
・文庫本感覚でお楽しみください

北海道の駅巡りコーナー

日高本線 車なし全駅巡り

ローカル線や秘境駅、マニアックな駅に興味のある方におすすめ! 2021年に大半区間が廃線になる、北海道の日高本線の全区間・全29駅(苫小牧~様似)を記録した本です。マイカーを使わずに、公共交通機関(バス)と徒歩のみで全駅訪問を行いました。日高本線が延伸する計画のあった、襟裳岬まで様似から足を伸ばしています。代行バスと路線バスの織り成す極限の時刻表ゲームと、絶海の太平洋と馬に囲まれた日高路、日高の隠れたグルメを是非たっぷり堪能してください。A4・フルカラー・192ページのたっぷりのボリュームで、あなたも旅行気分を漫喫できること待ったなし!

見どころ:日高本線被災区間(大狩部、慶能舞川橋梁、清畠~豊郷) / 牧場に囲まれた絵笛駅 / 窓口のあっただるま駅・荻伏駅 / 汐見の戦争遺跡のトーチカ / 新冠温泉、三石温泉 / 襟裳岬

A4 全192ページフルカラー / 2020年11月発行


Pocket
LINEで送る
Delicious にシェア

Add a Comment

メールアドレスが公開されることはありません。 * が付いている欄は必須項目です