TensorFlow 2がTensorFlow 1よりもはるかに遅いのはなぜですか？

Pytorchに切り替える理由として多くのユーザーから引用されていますが、私は熱心な実行のために最も重要な実用的な品質、速度を犠牲にする理由と説明をまだ見つけていません。

以下は、TF1とTF2を比較したコードベンチマークのパフォーマンスです。TF1は、47％から276％速く実行されています。

私の質問は、グラフまたはハードウェアレベルで、このような大幅な速度低下をもたらすのは何ですか？

詳細な答えを探しています-すでに幅広い概念に精通しています。関連Git

仕様：CUDA 10.0.130、cuDNN 7.4.2、Python 3.7.4、Windows 10、GTX 1070

ベンチマーク結果：

更新：以下のコードに従ってEager Executionを無効にしても効果はありません。ただし、動作に一貫性がありません。グラフモードで実行するとかなり役立つ場合もあれば、Eagerに比べて実行速度が遅い場合もあります。

TF開発者はどこにも現れないので、私はこの問題を自分で調査します。リンクされたGithubの問題の進捗状況を追跡できます。

更新2：説明に沿って、共有する大量の実験結果。今日行われるべきです。

ベンチマークコード：

# use tensorflow.keras... to benchmark tf.keras; used GPU for all above benchmarks
from keras.layers import Input, Dense, LSTM, Bidirectional, Conv1D
from keras.layers import Flatten, Dropout
from keras.models import Model
from keras.optimizers import Adam
import keras.backend as K
import numpy as np
from time import time

batch_shape = (32, 400, 16)
X, y = make_data(batch_shape)

model_small = make_small_model(batch_shape)
model_small.train_on_batch(X, y)  # skip first iteration which builds graph
timeit(model_small.train_on_batch, 200, X, y)

K.clear_session()  # in my testing, kernel was restarted instead

model_medium = make_medium_model(batch_shape)
model_medium.train_on_batch(X, y)  # skip first iteration which builds graph
timeit(model_medium.train_on_batch, 10, X, y)

使用した機能：

def timeit(func, iterations, *args):
    t0 = time()
    for _ in range(iterations):
        func(*args)
    print("Time/iter: %.4f sec" % ((time() - t0) / iterations))

def make_small_model(batch_shape):
    ipt   = Input(batch_shape=batch_shape)
    x     = Conv1D(128, 400, strides=4, padding='same')(ipt)
    x     = Flatten()(x)
    x     = Dropout(0.5)(x)
    x     = Dense(64, activation='relu')(x)
    out   = Dense(1,  activation='sigmoid')(x)
    model = Model(ipt, out)
    model.compile(Adam(lr=1e-4), 'binary_crossentropy')
    return model

def make_medium_model(batch_shape):
    ipt   = Input(batch_shape=batch_shape)
    x     = Bidirectional(LSTM(512, activation='relu', return_sequences=True))(ipt)
    x     = LSTM(512, activation='relu', return_sequences=True)(x)
    x     = Conv1D(128, 400, strides=4, padding='same')(x)
    x     = Flatten()(x)
    x     = Dense(256, activation='relu')(x)
    x     = Dropout(0.5)(x)
    x     = Dense(128, activation='relu')(x)
    x     = Dense(64,  activation='relu')(x)
    out   = Dense(1,   activation='sigmoid')(x)
    model = Model(ipt, out)
    model.compile(Adam(lr=1e-4), 'binary_crossentropy')
    return model

def make_data(batch_shape):
    return np.random.randn(*batch_shape), np.random.randint(0, 2, (batch_shape[0], 1))

UPDATE 2/18/2020：私は2.1と2.1-nightlyをベンチマークしました。結果はまちまちです。1つを除くすべての構成（モデルとデータサイズ）は、最高のTF2とTF1と同じかそれよりはるかに高速です。より遅く、劇的に遅くなるのは、Large-Large-espです。グラフ実行時（1.6xから2.5x遅い）。

さらに、私がテストした大きなモデルでは、グラフとイーガーの間に極端な再現性の違いがあります-ランダム性/並列計算では説明できないものです。現在、これらの時間あたりのクレームの再現可能なコードを提示することはできません。そのため、独自のモデルでこれをテストすることを強くお勧めします。

これらについてはまだGitの問題を開いていませんが、元のコメントはありました-まだ返答はありません。進捗状況がわかり次第、回答を更新します。

VERDICT：何をしているのかわかっていれば、そうではありません。しかし、そうしないと、平均で数個のGPUアップグレードによって、そして最悪の場合は複数のGPUによって、コストがかかる可能性があります。

THIS ANSWER：問題の概要と、ニーズに固有のトレーニング構成を決定する方法のガイドラインを提供することを目的としています。すべてのベンチマーク結果と使用したコードを含む詳細な低レベルの説明については、他の回答を参照してください。

回答がありましたら、詳細を確認しながら更新します。参考のために、この質問をブックマーク/「スター」できます。

問題の概要：TensorFlow開発者であるQ. Scott Zhuによって確認されたように、TF2はTFソー​​スの抜本的な変更（グラフレベルを含む）を伴う、Easerの実行とKerasとの緊密な統合に焦点を当てた開発を行いました。利点：大幅に拡張された処理、配布、デバッグ、および展開機能。ただし、これらのいくつかのコストは速度です。

ただし、問題はかなり複雑です。これはTF1とTF2だけではありません-列車の速度に大きな違いをもたらす要因は次のとおりです。

1. TF2対TF1
2. 熱心対グラフモード
3. keras 対 tf.keras
4. numpy対tf.data.Dataset対...
5. train_on_batch() 対 fit()
6. GPUとCPU
7. model(x)対model.predict(x)対...

残念ながら、上記のほとんどは互いに独立しておらず、それぞれが少なくとも実行時間を他のものと比べて2倍にすることができます。さいわい、いくつかのショートカットを使用して、体系的に最も効果的に機能するものを決定できます。

私は何をすべきか？現在、唯一の方法は、特定のモデル、データ、ハードウェアを試すことです。単一のコンフィギュレーションは、常に最適に動作しません-しかし、そこにあるやるのは、検索を簡素化するためだはありません。

>>する：

• train_on_batch()+ numpy+ tf.keras+ TF1 +熱心な/グラフ
• train_on_batch()+ numpy+ tf.keras+ TF2 +グラフ
• fit()+ numpy+ tf.keras+ TF1 / TF2 +グラフ+大きなモデルとデータ

>>してはいけないこと：

• fit()+ numpy+ keras中小規模のモデルとデータ
• fit()+ numpy+ tf.keras+ TF1 / TF2 +熱心な

• train_on_batch()+ numpy+ keras+ TF1 +熱心

• [少佐] tf.python.keras ; 実行速度は10〜100倍遅く、バグが多数あります。より詳しい情報

  • これにはlayers、models、optimizers「アウト・オブ・ボックスの」使用輸入関連、＆​​; ops、utils、および関連する「プライベート」インポートは問題ありませんが、確実に、altsを確認し、それらがtf.keras

ベンチマークの設定例については、他の回答の下部にあるコードを参照してください。上記のリストは、主に他の回答の「ベンチマーク」表に基づいています。

上記の禁止事項と禁止事項の制限：

• この質問のタイトルは「TF2がTF1よりもはるかに遅いのはなぜですか？」であり、その本体は明示的にトレーニングに関係していますが、問題はそれに限定されません。推論も、同じTFバージョン、インポート、データ形式などの中でさえ、大きな速度の違いの影響を受けます- この回答を参照してください。
• RNNはTF2で改善されているため、他の回答のデータグリッドを著しく変更する可能性があります。
• 主に使用されるモデルConv1DとDense-ないのRNN、疎データ/ターゲット、4 / 5Dの入力、および他のコンフィグ
• 入力データはnumpyおよびtf.data.Datasetに限定されますが、他の多くの形式が存在します。他の答えを見る
• GPUが使用されました。結果は CPUによって異なります。実際、質問したところ、CUDAが適切に構成されておらず、結果の一部はCPUベースでした。

TF2が熱心な実行のために、最も実用的な品質、速度を犠牲にしたのはなぜですか？明らかに、そうではありません-グラフはまだ利用可能です。しかし、質問が「なぜ熱心なのか」である場合：

• 優れたデバッグ：「中間層の出力を取得する方法」または「重みを検査する方法」を尋ねる多数の質問におそらく出くわします。熱心な場合、それは（ほぼ）と同じくらい簡単.__dict__です。対照的に、グラフは特別なバックエンド機能に精通している必要があります-デバッグとイントロスペクションのプロセス全体を非常に複雑にします。
• より高速なプロトタイピング：上記と同様のアイデアによる。より速く理解する=実際のDLのための残り時間。

EAGERを有効/無効にする方法は？

tf.enable_eager_execution()  # TF1; must be done before any model/tensor creation
tf.compat.v1.disable_eager_execution() # TF2; above holds

追加情報：

• _on_batch()TF2のメソッドに注意してください。TF開発者によると、彼らはまだ遅い実装を使用していますが、意図的にではありません。つまり、修正される予定です。詳細については、他の回答を参照してください。

TENSORFLOW DEVSへの要求：

1. を修正しtrain_on_batch()、fit()繰り返し呼び出すことのパフォーマンスの側面を修正してください。カスタムトレインループは、多くの人にとって、特に私にとって重要です。
2. ユーザーが理解できるように、これらのパフォーマンスの違いに関するドキュメント/ docstringの言及を追加します。
3. 全体的な実行速度を改善して、のぞき見をPytorchにホップしないようにします。

謝辞：ありがとう

• Q. 問題の詳細を説明してくれたTensorFlow開発者のScott Zhu氏。
• 有用なテストを共有してくれたP. Andrey と議論。

アップデート：

• 11 /14/ 19 - Numpy入力データを含むすべての*構成の TF2で実行が遅くなるモデル（私の実際のアプリケーション）を見つけました。差は13〜19％で、平均は17％でした。ただし、kerasとの違いtf.kerasはより劇的でした：18-40％、平均。32％（TF1と2の両方）。（*-TF2がOOMしたEagerを除く）

• 11/17/19-開発者on_batch()が最近のコミットでメソッドを更新し、速度が向上したと述べている-TF 2.1でリリースされるか、現在として入手可能tf-nightly。私は後者を実行することができないので、2.1までベンチを遅らせます。

• 2/20/20-予測パフォーマンスもベンチマークの価値があります。たとえば、TF2では、CPU予測時間に定期的なスパイクが含まれる可能性があります

THIS ANSWER：TF2対TF1トレインループ、入力データプロセッサ、Eager対グラフモードの実行など、問題の詳細なグラフ/ハードウェアレベルの説明を提供することを目的としています。問題の概要と解決のガイドラインについては、他の回答を参照してください。

パフォーマンスVERDICT：構成によっては、一方がより高速になることもあれば、もう一方が高速になることもあります。TF2とTF1に関する限り、それらは平均でほぼ同等ですが、重要な構成ベースの違いが存在し、TF1はTF2よりも頻繁にTF2よりも優先されます。以下の「ベンチマーク」を参照してください。

EAGER VS. グラフ：一部の人にとって、この全体の答えの要点：私のテストによると、TF2の熱意はTF1 よりも遅いです。詳細は下にあります。

この2つの基本的な違いは次のとおりです。グラフは計算ネットワークをプロアクティブにセットアップし、「指示された」ときに実行されます。一方、イーガーは作成時にすべてを実行します。しかし、話はここから始まります：

• EagerはGraphを欠いていないわけではなく、実際には、期待に反してほとんどが Graphである可能性があります。それが主に何であるか、実行されたグラフ -これは、グラフの大部分を構成するモデルとオプティマイザの重みを含みます。

• Eagerは実行時に自身のグラフの一部を再構築します。完全に構築されていないグラフの直接的な結果-プロファイラーの結果を参照してください。これには計算オーバーヘッドがあります。

• Numpy入力を使用すると、イーガーが遅くなります。このGitのコメントとコードに従って、EagerのNumpy入力には、テンソルをCPUからGPUにコピーするオーバーヘッドコストが含まれています。ソースコードをステップスルーすると、データ処理の違いは明らかです。EagerはNumpyを直接渡しますが、Graphはテンソルを渡します。テンソルはNumpyに評価されます。正確なプロセスは不明ですが、後者にはGPUレベルの最適化が必要です

• TF2 EagerはTF1 Eager よりも遅い -これは...予期せぬことです。以下のベンチマーク結果を参照してください。違いは無視できるものから重要なものまでありますが、一貫しています。理由が不明です-TF開発者が明確にした場合、回答が更新されます。

TF2対TF1：TF 開発者、Q。スコットジュー、応答の関連部分を引用します-私の強調と言い直しのw /

熱心に、ランタイムはopsを実行し、Pythonコードのすべての行の数値を返す必要があります。シングルステップ実行の性質により、処理が遅くなります。

TF2では、Kerasはtf.functionを利用して、トレーニング、評価、予測のためのグラフを作成します。それらをモデルの「実行関数」と呼びます。TF1では、「実行関数」はFuncGraphでしたが、TF関数としていくつかの共通コンポーネントを共有していますが、実装が異なります。

プロセス中に、train_on_batch（）、test_on_batch（）、およびpredict_on_batch（）の実装がどういうわけか残されました。これらは依然として数値的に正しいですが、x_on_batchの実行関数は、tf.functionでラップされたpython関数ではなく、純粋なpython関数です。これは遅くなります

TF2では、すべての入力データをtf.data.Datasetに変換します。これにより、実行関数を単一のタイプの入力を処理するように統合できます。データセットの変換にはある程度のオーバーヘッドがある可能性があります。これは、バッチごとのコストではなく、1回限りのオーバーヘッドであると思います

上記の最後の段落の最後の文と、下の段落の最後の句：

Eagerモードでの速度低下を克服するために、Python関数をグラフに変換する@ tf.functionがあります。np配列のような数値をフィードすると、tf.functionの本体が静的グラフに変換されて最適化され、最終値を返します。これは高速で、TF1グラフモードと同様のパフォーマンスを持つはずです。

同意しない-プロファイリングの結果によると、イーガーの入力データ処理はグラフの処理よりもかなり遅いことがわかります。また、tf.data.Dataset特に不明ですが、Eagerは同じデータ変換メソッドの複数を繰り返し呼び出します-プロファイラーを参照してください。

最後に、開発者のリンクされたコミット：Keras v2ループをサポートするための大幅な変更。

ループのトレーニング：（1）熱心対グラフに依存。（2）入力データ形式、トレーニングは異なるトレインループで進行します-TF2では_select_training_loop()、training.pyのいずれか：

training_v2.Loop()
training_distributed.DistributionMultiWorkerTrainingLoop(
              training_v2.Loop()) # multi-worker mode
# Case 1: distribution strategy
training_distributed.DistributionMultiWorkerTrainingLoop(
            training_distributed.DistributionSingleWorkerTrainingLoop())
# Case 2: generator-like. Input is Python generator, or Sequence object,
# or a non-distributed Dataset or iterator in eager execution.
training_generator.GeneratorOrSequenceTrainingLoop()
training_generator.EagerDatasetOrIteratorTrainingLoop()
# Case 3: Symbolic tensors or Numpy array-like. This includes Datasets and iterators 
# in graph mode (since they generate symbolic tensors).
training_generator.GeneratorLikeTrainingLoop() # Eager
training_arrays.ArrayLikeTrainingLoop() # Graph

それぞれが異なる方法でリソース割り当てを処理し、パフォーマンスと機能に影響を与えます。

Train Loops：fitvs train_on_batch、kerasvstf.keras .：4つのそれぞれが異なるトレインループを使用していますが、可能なすべての組み合わせではありません。keras" fit、例えば、フォームの使用fit_loop例を、training_arrays.fit_loop()と、そのはtrain_on_batch使用することができますK.function()。tf.keras前のセクションで一部説明したより洗練された階層があります。

ループのトレーニング：ドキュメント - いくつかの異なる実行方法に関する関連するソースdocstring：

他のTensorFlow演算とは異なり、Pythonの数値入力をテンソルに変換しません。さらに、Pythonの数値ごとに新しいグラフが生成されます

function 入力形状とデータ型の一意のセットごとに個別のグラフをインスタンス化します。

単一のtf.functionオブジェクトは、内部で複数の計算グラフにマップする必要がある場合があります。これはパフォーマンスとしてのみ表示されます（トレースグラフにはゼロ以外の計算コストとメモリコストがあります）

入力データプロセッサ：上記と同様に、ランタイム構成（実行モード、データ形式、配布戦略）に従って設定された内部フラグに応じて、ケースバイケースでプロセッサが選択されます。イージーで最も単純なケースは、Numpy配列で直接動作します。特定の例については、この回答を参照してください。

モデルサイズ、データサイズ：

• 決定的である; すべてのモデルとデータサイズの上位に立つ単一の構成はありません。
• モデルサイズに対するデータサイズは重要です。小さなデータとモデルの場合、データ転送（CPUからGPUなど）のオーバーヘッドが支配的になる可能性があります。同様に、小さなオーバーヘッドプロセッサは、データ変換時間が支配する大規模なデータでは、実行が遅くなる可能性があります（convert_to_tensor「プロファイラ」を参照 ）。
• 速度は、トレインループおよび入力データプロセッサのリソース処理方法によって異なります。

ベンチマーク：挽肉。- Word文書 - Excelスプレッドシート

用語：

• ％未満の数値はすべて秒です
• ％として計算(1 - longer_time / shorter_time)*100; 理論的根拠：私たちは、どちらがどちらがより速いかということに関心があります。shorter / longer実際には非線形関係であり、直接比較には役立ちません
• ％記号の決定：
  • TF2とTF1：+TF2の方が速い場合
  • GvE（グラフvsイーガー）：+グラフが高速の場合
• TF2 = TensorFlow 2.0.0 + Keras 2.3.1; TF1 = TensorFlow 1.14.0 + Keras 2.2.5

プロファイラー：

プロファイラ-説明：Spyder 3.3.6 IDEプロファイラ。

• 一部の機能は他の機能のネストで繰り返されます。したがって、「データ処理」関数と「トレーニング」関数の正確な分離を追跡するのは難しいため、一部のオーバーラップがあります-最後の結果で顕著です。

• ％計算されたwrtランタイムからビルド時間を引いた数値

• 1回または2回呼び出されたすべての（一意の）ランタイムを合計して計算されたビルド時間
• 反復回数と同じ回数呼び出されたすべての（一意の）ランタイム、およびそれらのネストのランタイムのいくつかを合計することによって計算されたトレーニング時間
• 関数は、残念ながら元の名前に従ってプロファイリングされ_func = funcます（つまり、としてプロファイリングされますfunc）。これは、ビルド時に混在するため、除外する必要があります。

テスト環境：

• 実行された最小限のバックグラウンドタスクで実行されたコード
• この投稿で提案されているように、GPUはタイミング反復の前に数回の反復で「ウォームアップ」されました
• ソースからビルドされたCUDA 10.0.130、cuDNN 7.6.0、TensorFlow 1.14.0、およびTensorFlow 2.0.0、さらにAnaconda
• Python 3.7.4、Spyder 3.3.6 IDE
• GTX 1070、Windows 10、24 GB DDR4 2.4 MHz RAM、i7-7700HQ 2.8 GHz CPU

方法論：

• 「小」、「中」、「大」のモデルとデータサイズのベンチマーク
• 入力データサイズに関係なく、各モデルサイズのパラメーター数を修正
• 「より大きな」モデルには、より多くのパラメーターとレイヤーがあります
• 「大きい」データはシーケンスが長くなりますが、同じでbatch_sizeあり、num_channels
• モデルConv1DはDense「学習可能な」レイヤーのみを使用します。RNNはTFバージョンの問題ごとに回避されました。違い
• モデルとオプティマイザグラフの構築を省略するために、ベンチマークループの外で常に1つのトレインフィットを実行しました
• スパースデータ（例layers.Embedding()：）またはスパースターゲット（例：）SparseCategoricalCrossEntropy()

制限：「完全な」答えは、可能なすべての列車ループと反復子を説明しますが、それは確かに私の時間能力、存在しない給与、または一般的な必要性を超えています。結果は方法論と同じくらい良いです-オープンマインドで解釈してください。

コード：

import numpy as np
import tensorflow as tf
import random
from termcolor import cprint
from time import time

from tensorflow.keras.layers import Input, Dense, Conv1D
from tensorflow.keras.layers import Dropout, GlobalAveragePooling1D
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
import tensorflow.keras.backend as K
#from keras.layers import Input, Dense, Conv1D
#from keras.layers import Dropout, GlobalAveragePooling1D
#from keras.models import Model 
#from keras.optimizers import Adam
#import keras.backend as K

#tf.compat.v1.disable_eager_execution()
#tf.enable_eager_execution()

def reset_seeds(reset_graph_with_backend=None, verbose=1):
    if reset_graph_with_backend is not None:
        K = reset_graph_with_backend
        K.clear_session()
        tf.compat.v1.reset_default_graph()
        if verbose:
            print("KERAS AND TENSORFLOW GRAPHS RESET")

    np.random.seed(1)
    random.seed(2)
    if tf.__version__[0] == '2':
        tf.random.set_seed(3)
    else:
        tf.set_random_seed(3)
    if verbose:
        print("RANDOM SEEDS RESET")

print("TF version: {}".format(tf.__version__))
reset_seeds()

def timeit(func, iterations, *args, _verbose=0, **kwargs):
    t0 = time()
    for _ in range(iterations):
        func(*args, **kwargs)
        print(end='.'*int(_verbose))
    print("Time/iter: %.4f sec" % ((time() - t0) / iterations))

def make_model_small(batch_shape):
    ipt   = Input(batch_shape=batch_shape)
    x     = Conv1D(128, 40, strides=4, padding='same')(ipt)
    x     = GlobalAveragePooling1D()(x)
    x     = Dropout(0.5)(x)
    x     = Dense(64, activation='relu')(x)
    out   = Dense(1,  activation='sigmoid')(x)
    model = Model(ipt, out)
    model.compile(Adam(lr=1e-4), 'binary_crossentropy')
    return model

def make_model_medium(batch_shape):
    ipt = Input(batch_shape=batch_shape)
    x = ipt
    for filters in [64, 128, 256, 256, 128, 64]:
        x  = Conv1D(filters, 20, strides=1, padding='valid')(x)
    x     = GlobalAveragePooling1D()(x)
    x     = Dense(256, activation='relu')(x)
    x     = Dropout(0.5)(x)
    x     = Dense(128, activation='relu')(x)
    x     = Dense(64,  activation='relu')(x)
    out   = Dense(1,   activation='sigmoid')(x)
    model = Model(ipt, out)
    model.compile(Adam(lr=1e-4), 'binary_crossentropy')
    return model

def make_model_large(batch_shape):
    ipt   = Input(batch_shape=batch_shape)
    x     = Conv1D(64,  400, strides=4, padding='valid')(ipt)
    x     = Conv1D(128, 200, strides=1, padding='valid')(x)
    for _ in range(40):
        x = Conv1D(256,  12, strides=1, padding='same')(x)
    x     = Conv1D(512,  20, strides=2, padding='valid')(x)
    x     = Conv1D(1028, 10, strides=2, padding='valid')(x)
    x     = Conv1D(256,   1, strides=1, padding='valid')(x)
    x     = GlobalAveragePooling1D()(x)
    x     = Dense(256, activation='relu')(x)
    x     = Dropout(0.5)(x)
    x     = Dense(128, activation='relu')(x)
    x     = Dense(64,  activation='relu')(x)    
    out   = Dense(1,   activation='sigmoid')(x)
    model = Model(ipt, out)
    model.compile(Adam(lr=1e-4), 'binary_crossentropy')
    return model

def make_data(batch_shape):
    return np.random.randn(*batch_shape), \
           np.random.randint(0, 2, (batch_shape[0], 1))

def make_data_tf(batch_shape, n_batches, iters):
    data = np.random.randn(n_batches, *batch_shape),
    trgt = np.random.randint(0, 2, (n_batches, batch_shape[0], 1))
    return tf.data.Dataset.from_tensor_slices((data, trgt))#.repeat(iters)

batch_shape_small  = (32, 140,   30)
batch_shape_medium = (32, 1400,  30)
batch_shape_large  = (32, 14000, 30)

batch_shapes = batch_shape_small, batch_shape_medium, batch_shape_large
make_model_fns = make_model_small, make_model_medium, make_model_large
iterations = [200, 100, 50]
shape_names = ["Small data",  "Medium data",  "Large data"]
model_names = ["Small model", "Medium model", "Large model"]

def test_all(fit=False, tf_dataset=False):
    for model_fn, model_name, iters in zip(make_model_fns, model_names, iterations):
        for batch_shape, shape_name in zip(batch_shapes, shape_names):
            if (model_fn is make_model_large) and (batch_shape is batch_shape_small):
                continue
            reset_seeds(reset_graph_with_backend=K)
            if tf_dataset:
                data = make_data_tf(batch_shape, iters, iters)
            else:
                data = make_data(batch_shape)
            model = model_fn(batch_shape)

            if fit:
                if tf_dataset:
                    model.train_on_batch(data.take(1))
                    t0 = time()
                    model.fit(data, steps_per_epoch=iters)
                    print("Time/iter: %.4f sec" % ((time() - t0) / iters))
                else:
                    model.train_on_batch(*data)
                    timeit(model.fit, iters, *data, _verbose=1, verbose=0)
            else:
                model.train_on_batch(*data)
                timeit(model.train_on_batch, iters, *data, _verbose=1)
            cprint(">> {}, {} done <<\n".format(model_name, shape_name), 'blue')
            del model

test_all(fit=True, tf_dataset=False)