Mengapa TensorFlow 2 lebih lambat dari TensorFlow 1?

Ini telah dikutip oleh banyak pengguna sebagai alasan untuk beralih ke Pytorch, tetapi saya belum menemukan alasan / penjelasan untuk mengorbankan kualitas praktis yang paling penting, kecepatan, untuk pelaksanaan yang bersemangat.

Di bawah ini adalah kinerja pembandingan kode, TF1 vs TF2 - dengan TF1 berjalan di mana saja dari 47% menjadi 276% lebih cepat .

Pertanyaan saya adalah: apakah itu, pada grafik atau tingkat perangkat keras, yang menghasilkan perlambatan yang signifikan?

Mencari jawaban terinci - saya sudah terbiasa dengan konsep luas. Git yang relevan

Spesifikasi : CUDA 10.0.130, cuDNN 7.4.2, Python 3.7.4, Windows 10, GTX 1070

Hasil patok banding :

UPDATE : Menonaktifkan Eksekusi Eager per kode di bawah ini tidak membantu. Perilaku ini, bagaimanapun, tidak konsisten: kadang-kadang berjalan dalam mode grafik sangat membantu, di lain waktu itu berjalan lebih lambat dibandingkan dengan Eager.

Karena TF devs tidak muncul di mana-mana, saya akan menyelidiki masalah ini sendiri - dapat mengikuti kemajuan dalam masalah Github yang tertaut.

PEMBARUAN 2 : banyak hasil percobaan untuk dibagikan, beserta penjelasannya; harus dilakukan hari ini.

Kode benchmark :

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

Fungsi yang digunakan :

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 : Saya telah duduk di bangku 2.1 dan 2.1 malam; hasilnya dicampur. Semua kecuali satu konfigurasi (model & ukuran data) secepat atau jauh lebih cepat daripada yang terbaik dari TF2 & TF1. Yang lebih lambat, dan lebih lambat secara dramatis, adalah Besar-Besar - esp. dalam eksekusi Grafik ( 1,6x ke 2,5x lebih lambat ).

Selain itu, ada perbedaan reproduktifitas ekstrim antara Graph dan Eager untuk model besar yang saya uji - yang tidak dapat dijelaskan melalui keacakan / kompute-paralelisme. Saat ini saya tidak dapat menyajikan kode yang dapat direproduksi untuk klaim ini per batasan waktu, jadi alih-alih saya sangat merekomendasikan pengujian ini untuk model Anda sendiri.

Belum membuka masalah Git tentang ini, tapi saya mengomentari aslinya - belum ada tanggapan. Saya akan memperbarui jawaban setelah kemajuan dibuat.

VERDICT : tidak , JIKA Anda tahu apa yang Anda lakukan. Tetapi jika Anda tidak melakukannya , harganya mungkin mahal, banyak - dengan beberapa peningkatan GPU rata-rata, dan oleh beberapa kasus terburuk GPU.

JAWABAN INI : bertujuan untuk memberikan deskripsi tingkat tinggi tentang masalah ini, serta pedoman untuk bagaimana memutuskan konfigurasi pelatihan khusus untuk kebutuhan Anda. Untuk deskripsi terperinci, tingkat rendah, yang mencakup semua hasil + kode pembandingan yang digunakan, lihat jawaban saya yang lain.

Saya akan memperbarui jawaban saya dengan info lebih lanjut jika saya mengetahui ada - dapat menandai / "membintangi" pertanyaan ini untuk referensi.

RINGKASAN MASALAH : sebagaimana dikonfirmasi oleh pengembang TensorFlow, Q. Scott Zhu, TF2 memfokuskan pengembangan pada pelaksanaan Eager & integrasi yang ketat dengan Keras, yang melibatkan perubahan menyeluruh pada sumber TF - termasuk di tingkat grafik. Manfaat: kemampuan pemrosesan, distribusi, debug, dan penyebaran yang sangat diperluas. Biaya beberapa di antaranya, bagaimanapun, adalah kecepatan.

Masalahnya, bagaimanapun, cukup rumit. Ini bukan hanya TF1 vs TF2 - faktor-faktor yang menghasilkan perbedaan kecepatan kereta yang signifikan termasuk:

1. TF2 vs TF1
2. Modus bersemangat vs. Grafik
3. keras vs. tf.keras
4. numpyvs. tf.data.Datasetvs. ...
5. train_on_batch() vs. fit()
6. GPU vs. CPU
7. model(x)vs. model.predict(x)vs. ...

Sayangnya, hampir tidak ada di atas yang independen dari yang lain, dan masing-masing setidaknya dapat menggandakan waktu eksekusi relatif terhadap yang lain. Untungnya, Anda dapat menentukan apa yang akan bekerja paling baik secara sistematis, dan dengan beberapa cara pintas - seperti yang akan saya tunjukkan.

APA YANG HARUS SAYA LAKUKAN? Saat ini, satu-satunya cara adalah - bereksperimen untuk model, data, dan perangkat keras spesifik Anda. Tidak ada satu konfigurasi yang akan selalu berfungsi dengan baik - tetapi ada yang harus dan yang tidak boleh dilakukan untuk menyederhanakan pencarian Anda:

>> LAKUKAN:

• train_on_batch()+ numpy+ tf.keras+ TF1 + Eager / Grafik
• train_on_batch()+ numpy+ tf.keras+ TF2 + Grafik
• fit()+ numpy+ tf.keras+ TF1 / TF2 + Grafik + model & data besar

>> JANGAN:

• fit()+ numpy+ kerasUntuk model kecil & menengah dan data
• fit()+ numpy+ tf.keras+ TF1 / TF2 + Bersemangat

• train_on_batch()+ numpy+ keras+ TF1 + Ingin

• [Mayor] tf.python.keras ; itu bisa berjalan 10-100x lebih lambat, dan dengan banyak bug; Info lebih lanjut

  • Ini termasuk layers, models, optimizers, & terkait "out-of-box" impor penggunaan; ops, utils, & impor 'pribadi' terkait baik-baik saja - tetapi untuk memastikan, periksa alts, & apakah mereka digunakan ditf.keras

Lihat kode di bagian bawah jawaban saya yang lain untuk contoh pengaturan benchmarking. Daftar di atas didasarkan terutama pada tabel "BENCHMARKS" di jawaban lain.

PEMBATASAN atas DO & JANGAN di atas:

• Pertanyaan ini berjudul "Mengapa TF2 jauh lebih lambat daripada TF1?", Dan sementara tubuhnya menyangkut pelatihan secara eksplisit, masalah ini tidak terbatas pada itu; inferensi , juga tunduk pada perbedaan kecepatan utama, bahkan dalam versi TF yang sama, impor, format data, dll. - lihat jawaban ini .
• RNN cenderung mengubah kisi data di jawaban lain, karena mereka telah ditingkatkan di TF2
• Model terutama digunakan Conv1Ddan Dense- tidak ada RNN, data / target jarang, input 4 / 5D, & konfigurasi lainnya
• Input data terbatas pada numpydan tf.data.Dataset, sementara banyak format lain ada; lihat jawaban lain
• GPU digunakan; hasilnya akan berbeda pada CPU. Bahkan, ketika saya mengajukan pertanyaan, CUDA saya tidak dikonfigurasi dengan benar, dan beberapa hasilnya berbasis CPU.

Mengapa TF2 mengorbankan kualitas, kecepatan, untuk eksekusi yang paling praktis? Tidak, jelas - grafik masih tersedia. Tetapi jika pertanyaannya adalah "mengapa bersemangat sekali":

• Debugging unggul : Anda mungkin menemukan banyak pertanyaan yang menanyakan "bagaimana cara mendapatkan output lapisan menengah" atau "bagaimana cara memeriksa bobot"; dengan penuh semangat, itu (hampir) sesederhana .__dict__. Grafik, sebaliknya, membutuhkan keakraban dengan fungsi backend khusus - sangat menyulitkan seluruh proses debugging & introspeksi.
• Pembuatan prototipe lebih cepat : per ide yang mirip dengan di atas; pemahaman yang lebih cepat = lebih banyak waktu tersisa untuk DL aktual.

BAGAIMANA CARA MENGaktifkan / MENONAKTIFKAN EAGER?

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

INFO TAMBAHAN :

• Hati-hati dengan _on_batch()metode dalam TF2; menurut TF dev, mereka masih menggunakan implementasi lebih lambat, tetapi tidak sengaja - yaitu itu harus diperbaiki. Lihat jawaban lain untuk detailnya.

PERMINTAAN UNTUK PERANGKAT TENSORFLOW :

1. Harap perbaiki train_on_batch(), dan aspek kinerja dari panggilan fit()iteratif; loop kereta adat penting bagi banyak orang, terutama bagi saya.
2. Tambahkan dokumentasi / penyebutan penyebutan perbedaan kinerja ini untuk pengetahuan pengguna.
3. Tingkatkan kecepatan eksekusi umum agar peeps tidak melompat ke Pytorch.

UCAPAN TERIMA KASIH : Terima kasih kepada

• P. Scott Zhu, pengembang TensorFlow, untuk klarifikasi terperinci tentang masalah ini.
• P. Andrey untuk berbagi pengujian dan diskusi yang bermanfaat.

PEMBARUAN :

• 11/14/19 - menemukan model (dalam aplikasi saya yang sebenarnya) yang berjalan lebih lambat pada TF2 untuk semua * konfigurasi dengan data input Numpy. Perbedaan berkisar 13-19%, rata-rata 17%. Perbedaan antara kerasdan tf.keras, bagaimanapun, lebih dramatis: 18-40% , rata-rata. 32% (baik TF1 & 2). (* - kecuali Eager, di mana TF2 OOM)

• 11/17/19 - devs memperbarui on_batch()metode dalam komit baru - baru ini , yang menyatakan telah meningkatkan kecepatan - untuk dirilis dalam TF 2.1, atau tersedia sekarang sebagai tf-nightly. Karena saya tidak bisa menjalankan yang terakhir, akan menunda bangku hingga 2.1.

• 2/20/20 - kinerja prediksi juga layak dilakukan; dalam TF2, misalnya, waktu prediksi CPU dapat melibatkan lonjakan periodik

JAWABAN INI : bertujuan untuk memberikan deskripsi rinci, grafik / tingkat perangkat keras masalah - termasuk loop kereta TF2 vs TF1, prosesor data input, dan eksekusi mode Eager vs. Graph. Untuk ringkasan masalah & pedoman resolusi, lihat jawaban saya yang lain.

VERDICT PERFORMANCE : kadang-kadang satu lebih cepat, kadang-kadang yang lain, tergantung pada konfigurasi. Sejauh TF2 vs TF1 berjalan, mereka hampir setara, tetapi perbedaan berbasis konfigurasi yang signifikan memang ada, dan TF1 mengalahkan TF2 lebih sering daripada sebaliknya. Lihat "BENCHMARKING" di bawah ini.

EAGER VS. GRAFIK : daging dari seluruh jawaban ini untuk beberapa: keinginan TF2 lebih lambat dari pada TF1, menurut pengujian saya. Detail lebih jauh ke bawah.

Perbedaan mendasar antara keduanya adalah: Grafik mengatur jaringan komputasi secara proaktif , dan dijalankan ketika 'disuruh' - sedangkan Eager mengeksekusi semuanya saat penciptaan. Tetapi cerita hanya dimulai di sini:

• Bersemangat TIDAK tanpa Grafik , dan mungkin sebagian besar Grafik, bertentangan dengan harapan. Apa itu sebagian besar, dijalankan Grafik - ini termasuk bobot model & pengoptimal, terdiri dari sebagian besar grafik.

• Eager membangun kembali sebagian dari grafiknya sendiri pada saat eksekusi ; konsekuensi langsung dari Grafik yang tidak sepenuhnya dibangun - lihat hasil profiler. Ini memiliki overhead komputasi.

• Eager lebih lambat dengan input Numpy ; sesuai komentar & kode Git ini , input Numpy di Eager termasuk biaya overhead untuk menyalin tensor dari CPU ke GPU. Melangkah melalui kode sumber, perbedaan penanganan data jelas; Bersemangat langsung melewati Numpy, sementara Grafik melewati tensor yang kemudian mengevaluasi ke Numpy; tidak yakin proses pastinya, tetapi yang terakhir harus melibatkan optimasi level GPU

• TF2 Eager lebih lambat dari TF1 Eager - ini ... tidak terduga. Lihat hasil pembandingan di bawah ini. Perbedaan rentang dari diabaikan hingga signifikan, tetapi konsisten. Tidak yakin mengapa demikian - jika seorang TF mengklarifikasi, akan memperbarui jawaban.

TF2 vs. TF1 : mengutip bagian yang relevan dari TF dev, Q. Scott Zhu's, respons - dengan sedikit penekanan & penulisan ulang saya:

Dalam keinginan, runtime perlu menjalankan ops dan mengembalikan nilai numerik untuk setiap baris kode python. Sifat eksekusi satu langkah menyebabkannya menjadi lambat .

Dalam TF2, Keras memanfaatkan fungsi untuk membuat grafiknya untuk pelatihan, evaluasi dan prediksi. Kami menyebutnya "fungsi eksekusi" untuk model. Dalam TF1, "fungsi eksekusi" adalah FuncGraph, yang berbagi beberapa komponen umum sebagai fungsi TF, tetapi memiliki implementasi yang berbeda.

Selama proses, kami entah bagaimana meninggalkan implementasi yang salah untuk train_on_batch (), test_on_batch () dan predict_on_batch () . Mereka masih benar secara numerik , tetapi fungsi eksekusi untuk x_on_batch adalah fungsi python murni, alih-alih fungsi python dibungkus tf.fungsi. Ini akan menyebabkan kelambatan

Dalam TF2, kami mengubah semua data input menjadi tf.data.Dataset, yang dengannya kami dapat menyatukan fungsi eksekusi kami untuk menangani jenis input tunggal. Mungkin ada beberapa overhead dalam konversi dataset , dan saya pikir ini hanya overhead satu kali, daripada biaya per-batch

Dengan kalimat terakhir dari paragraf terakhir di atas, dan klausa terakhir dari paragraf di bawah:

Untuk mengatasi kelambatan dalam mode eager, kita memiliki fungsi @tf., yang akan mengubah fungsi python menjadi grafik. Ketika mengumpankan nilai numerik seperti array np, fungsi tf.fungsi diubah menjadi grafik statis, dioptimalkan, dan mengembalikan nilai akhir, yang cepat dan harus memiliki kinerja yang sama dengan mode grafik TF1.

Saya tidak setuju - per hasil profil saya, yang menunjukkan pemrosesan data input Eager jauh lebih lambat daripada Graph. Juga, tidak yakin tentang tf.data.Datasetkhususnya, tetapi Eager berulang kali memanggil beberapa metode konversi data yang sama - lihat profiler.

Terakhir, komitmen tertaut dev: Jumlah perubahan signifikan untuk mendukung loop Keras v2 .

Train Loops : tergantung pada (1) Eager vs. Graph; (2) memasukkan format data, pelatihan akan dilanjutkan dengan loop kereta yang berbeda - di TF2 _select_training_loop(),, training.py , salah satu dari:

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

Masing-masing menangani alokasi sumber daya secara berbeda, dan membawa konsekuensi pada kinerja & kemampuan.

Train Loops: fitvs train_on_batch, kerasvstf.keras .: masing-masing dari empat menggunakan loop kereta yang berbeda, meskipun mungkin tidak dalam setiap kombinasi yang memungkinkan. keras' fit, misalnya, menggunakan bentuk fit_loop, misalnya training_arrays.fit_loop(), dan train_on_batchmungkin menggunakan K.function(). tf.kerasmemiliki hierarki yang lebih canggih yang dijelaskan pada bagian di bagian sebelumnya.

Train Loops: dokumentasi - dokumentasi sumber yang relevan tentang beberapa metode eksekusi yang berbeda:

Tidak seperti operasi TensorFlow lainnya, kami tidak mengonversi input numerik python ke tensor. Selain itu, grafik baru dihasilkan untuk setiap nilai numerik python yang berbeda

function instantiate grafik terpisah untuk setiap set unik bentuk input dan tipe data .

Objek fungsi tf. tunggal mungkin perlu memetakan ke beberapa grafik perhitungan di bawah tenda. Ini harus terlihat hanya sebagai kinerja (melacak grafik memiliki biaya komputasi dan memori bukan nol )

Input data processor : mirip dengan di atas, prosesor dipilih kasus per kasus, tergantung pada flag internal yang diatur sesuai dengan konfigurasi runtime (mode eksekusi, format data, strategi distribusi). Kasing paling sederhana dengan Eager, yang bekerja langsung dengan array Numpy. Untuk beberapa contoh spesifik, lihat jawaban ini .

UKURAN MODEL, UKURAN DATA:

• Sangat menentukan; tidak ada konfigurasi tunggal dimahkotai sendiri di atas semua model & ukuran data.
• Ukuran data relatif terhadap ukuran model adalah penting; untuk data & model kecil, transfer data (mis. CPU ke GPU) dapat mendominasi. Demikian juga, prosesor overhead kecil dapat berjalan lebih lambat pada data besar per mendominasi waktu konversi data (lihat convert_to_tensordi "PROFILER")
• Kecepatan berbeda per loop berbeda dan sarana pengolah data yang berbeda dalam menangani sumber daya.

BENCHMARKS : daging yang digiling. - Dokumen Word - Excel Spreadsheet

Terminologi :

• % -jumlah angka semua detik
• % dihitung sebagai (1 - longer_time / shorter_time)*100; alasan: kami tertarik dengan faktor apa yang lebih cepat dari yang lain; shorter / longersebenarnya adalah hubungan non-linear, tidak berguna untuk perbandingan langsung
• % penentuan tanda:
  • TF2 vs TF1: +jika TF2 lebih cepat
  • GvE (Grafik vs. Eager): +jika Grafik lebih cepat
• TF2 = TensorFlow 2.0.0 + Keras 2.3.1; TF1 = TensorFlow 1.14.0 + Keras 2.2.5

PROFILER :

PROFILER - Penjelasan : Spyder 3.3.6 IDE profiler.

• Beberapa fungsi diulang di sarang orang lain; karenanya, sulit untuk melacak pemisahan yang tepat antara fungsi "pemrosesan data" dan "pelatihan", sehingga akan ada beberapa tumpang tindih - seperti yang diucapkan dalam hasil terakhir.

• % angka dihitung waktu kerja dikurangi dikurangi waktu pembuatan

• Build time dihitung dengan menjumlahkan semua runtime (unik) yang disebut 1 atau 2 kali
• Latih waktu dihitung dengan menjumlahkan semua runtime (unik) yang disebut # kali sama dengan # iterasi, dan beberapa runtime sarangnya
• Fungsi diprofilkan sesuai dengan nama aslinya , sayangnya (yaitu _func = funcakan diprofilkan sebagai func), yang bercampur dalam waktu pembuatan - karenanya perlu untuk mengecualikannya

LINGKUNGAN PENGUJIAN :

• Kode yang dijalankan di bawah dengan tugas latar belakang minimal sedang berjalan
• GPU "dihangatkan" dengan beberapa iterasi sebelum waktu iterasi, seperti yang disarankan dalam posting ini
• CUDA 10.0.130, cuDNN 7.6.0, TensorFlow 1.14.0, & TensorFlow 2.0.0 dibangun dari sumber, ditambah Anaconda
• Python 3.7.4, Spyder 3.3.6 IDE
• GTX 1070, Windows 10, 24GB DDR4 2.4-MHz RAM, i7-7700HQ 2.8-GHz CPU

METODOLOGI :

• Ukuran & model data & ukuran 'kecil', 'sedang', & 'besar'
• Perbaiki # parameter untuk setiap ukuran model, terlepas dari ukuran data input
• Model "Lebih besar" memiliki lebih banyak parameter dan lapisan
• "Lebih besar" data memiliki urutan yang lebih panjang, tetapi sama batch_sizedannum_channels
• Model hanya menggunakan Conv1D, Dense'dipelajari' lapisan; RNN dihindari per implem versi TF. perbedaan
• Selalu jalankan satu train fit di luar loop pembandingan, untuk menghilangkan model & pembuatan grafik pengoptimal
• Tidak menggunakan data jarang (mis. layers.Embedding()) Atau target jarang (misSparseCategoricalCrossEntropy()

BATASAN : jawaban "lengkap" akan menjelaskan setiap loop kereta & iterator yang mungkin, tetapi itu jelas di luar kemampuan waktu saya, gaji tidak ada, atau kebutuhan umum. Hasilnya hanya sebagus metodologi - menafsirkan dengan pikiran terbuka.

Kode :

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)