SavedModels dari TF Hub di TensorFlow 2

Format SavedModel dari TensorFlow 2 adalah cara yang disarankan untuk membagikan model dan potongan model yang telah dilatih sebelumnya di TensorFlow Hub. Ini menggantikan format TF1 Hub yang lebih lama dan dilengkapi dengan serangkaian API baru.

Halaman ini menjelaskan cara menggunakan kembali TF2 SavedModels dalam program TensorFlow 2 dengan hub.load() API tingkat rendah dan pembungkus hub.KerasLayer -nya. (Biasanya, hub.KerasLayer digabungkan dengan tf.keras.layers lain untuk membuat model Keras atau model_fn dari Penaksir TF2.) API ini juga dapat memuat model lama dalam format Hub TF1, dalam batasan, lihat panduan kompatibilitas .

Pengguna TensorFlow 1 dapat memperbarui ke TF 1.15 dan kemudian menggunakan API yang sama. Versi TF1 yang lebih lama tidak berfungsi.

Menggunakan SavedModels dari TF Hub

Menggunakan SavedModel di Keras

Keras adalah API tingkat tinggi TensorFlow untuk membangun model pembelajaran mendalam dengan menyusun objek Lapisan Keras. Pustaka tensorflow_hub menyediakan hub.KerasLayer yang diinisialisasi dengan URL (atau jalur sistem file) dari SavedModel dan kemudian menyediakan komputasi dari SavedModel, termasuk bobot yang telah dilatih sebelumnya.

Berikut adalah contoh penggunaan penyematan teks yang telah dilatih sebelumnya:

import tensorflow as tf
import tensorflow_hub as hub

hub_url = "https://tfhub.dev/google/nnlm-en-dim128/2"
embed = hub.KerasLayer(hub_url)
embeddings = embed(["A long sentence.", "single-word", "http://example.com"])
print(embeddings.shape, embeddings.dtype)

Dari sini, pengklasifikasi teks dapat dibangun dengan cara Keras biasa:

model = tf.keras.Sequential([
    embed,
    tf.keras.layers.Dense(16, activation="relu"),
    tf.keras.layers.Dense(1, activation="sigmoid"),
])

Colab klasifikasi Teks adalah contoh lengkap cara melatih dan mengevaluasi pengklasifikasi semacam itu.

Bobot model di hub.KerasLayer diatur ke non-trainable secara default. Lihat bagian fine-tuning di bawah ini untuk mengetahui cara mengubahnya. Bobot dibagi antara semua aplikasi dari objek lapisan yang sama, seperti biasa di Keras.

Menggunakan Model Tersimpan di Penaksir

Pengguna API Estimator TensorFlow untuk pelatihan terdistribusi dapat menggunakan SavedModels dari TF Hub dengan menulis model_fn mereka dalam bentuk hub.KerasLayer di antara tf.keras.layers lainnya .

Di balik layar: Pengunduhan dan penyimpanan Model Tersimpan

Menggunakan SavedModel dari TensorFlow Hub (atau server HTTPS lain yang menerapkan protokol hostingnya ) mengunduh dan mendekompresnya ke sistem file lokal jika belum ada. Variabel lingkungan TFHUB_CACHE_DIR dapat diatur untuk mengganti lokasi sementara default untuk menyimpan cache dari SavedModels yang diunduh dan tidak dikompresi. Untuk detailnya, lihat Caching .

Menggunakan SavedModel di TensorFlow tingkat rendah

Fungsi hub.load(handle) mendownload dan mendekompresi SavedModel (kecuali handle sudah menjadi jalur sistem file) lalu mengembalikan hasil pemuatannya dengan fungsi tf.saved_model.load() . Oleh karena itu, hub.load() dapat menangani SavedModel yang valid (tidak seperti hub.Module pendahulunya untuk TF1).

Topik lanjutan: apa yang diharapkan dari SavedModel setelah memuat

Bergantung pada konten SavedModel, hasil dari obj = hub.load(...) dapat dipanggil dengan berbagai cara (seperti yang dijelaskan lebih detail di TensorFlow's SavedModel Guide :

  • Tanda tangan penyajian dari SavedModel (jika ada) direpresentasikan sebagai kamus fungsi konkret dan dapat disebut seperti tensors_out = obj.signatures["serving_default"](**tensors_in) , dengan kamus tensor yang dikunci oleh masing-masing input dan output nama dan tunduk pada batasan bentuk dan tipe tanda tangan.

  • Metode @tf.function dari objek yang disimpan (jika ada) dipulihkan sebagai objek tf.function yang dapat dipanggil oleh semua kombinasi argumen Tensor dan non-Tensor yang tf.function telah dilacak sebelum disimpan. Secara khusus, jika ada metode obj.__call__ dengan jejak yang sesuai, obj itu sendiri dapat dipanggil seperti fungsi Python. Contoh sederhana bisa terlihat seperti output_tensor = obj(input_tensor, training=False) .

Ini meninggalkan kebebasan besar dalam antarmuka yang dapat diimplementasikan oleh SavedModels. Antarmuka SavedModels yang Dapat Digunakan Kembali untuk obj menetapkan konvensi sedemikian rupa sehingga kode klien, termasuk adaptor seperti hub.KerasLayer , tahu cara menggunakan SavedModel.

Beberapa SavedModels mungkin tidak mengikuti konvensi itu, terutama seluruh model yang tidak dimaksudkan untuk digunakan kembali dalam model yang lebih besar, dan hanya memberikan tanda tangan penyajian.

Variabel yang dapat dilatih dalam SavedModel dimuat ulang sebagai dapat dilatih, dan tf.GradientTape akan menontonnya secara default. Lihat bagian fine-tuning di bawah ini untuk beberapa peringatan, dan pertimbangkan untuk menghindari ini sebagai permulaan. Bahkan jika Anda ingin menyempurnakan, Anda mungkin ingin melihat apakah obj.trainable_variables menyarankan untuk melatih ulang hanya sebagian dari variabel yang awalnya dapat dilatih.

Membuat Model Tersimpan untuk TF Hub

Ringkasan

SavedModel adalah format serialisasi standar TensorFlow untuk model terlatih atau potongan model. Ini menyimpan bobot terlatih model bersama dengan operasi TensorFlow yang tepat untuk melakukan perhitungannya. Itu dapat digunakan secara independen dari kode yang membuatnya. Secara khusus, ini dapat digunakan kembali di berbagai API pembuatan model tingkat tinggi seperti Keras, karena operasi TensorFlow adalah bahasa dasarnya yang umum.

Menyimpan dari Keras

Dimulai dengan TensorFlow 2, tf.keras.Model.save() dan tf.keras.models.save_model() default ke format SavedModel (bukan HDF5). SavedModels yang dihasilkan yang dapat digunakan dengan hub.load() , hub.KerasLayer dan adaptor serupa untuk API tingkat tinggi lainnya saat tersedia.

Untuk membagikan Model Keras yang lengkap, simpan saja dengan include_optimizer=False .

Untuk membagikan potongan Model Keras, buat potongan Model itu sendiri dan simpan. Anda dapat meletakkan kode seperti itu dari awal....

piece_to_share = tf.keras.Model(...)
full_model = tf.keras.Sequential([piece_to_share, ...])
full_model.fit(...)
piece_to_share.save(...)

...atau potong bagian untuk dibagikan setelah fakta (jika sesuai dengan layering model lengkap Anda):

full_model = tf.keras.Model(...)
sharing_input = full_model.get_layer(...).get_output_at(0)
sharing_output = full_model.get_layer(...).get_output_at(0)
piece_to_share = tf.keras.Model(sharing_input, sharing_output)
piece_to_share.save(..., include_optimizer=False)

Model TensorFlow di GitHub menggunakan pendekatan sebelumnya untuk BERT (lihat nlp/tools/export_tfhub_lib.py , perhatikan pemisahan antara core_model untuk ekspor dan pretrainer untuk memulihkan pos pemeriksaan) dan pendekatan terakhir untuk ResNet (lihat legacy/image_classification/tfhub_export.py ).

Menyimpan dari TensorFlow tingkat rendah

Ini membutuhkan pemahaman yang baik tentang Panduan SavedModel TensorFlow .

Jika Anda ingin memberikan lebih dari sekadar tanda tangan penyajian, Anda harus mengimplementasikan antarmuka Reusable SavedModel . Secara konseptual, ini terlihat seperti

class MyMulModel(tf.train.Checkpoint):
  def __init__(self, v_init):
    super().__init__()
    self.v = tf.Variable(v_init)
    self.variables = [self.v]
    self.trainable_variables = [self.v]
    self.regularization_losses = [
        tf.function(input_signature=[])(lambda: 0.001 * self.v**2),
    ]

  @tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
  def __call__(self, inputs):
    return tf.multiply(inputs, self.v)

tf.saved_model.save(MyMulModel(2.0), "/tmp/my_mul")

layer = hub.KerasLayer("/tmp/my_mul")
print(layer([10., 20.]))  # [20., 40.]
layer.trainable = True
print(layer.trainable_weights)  # [2.]
print(layer.losses)  # 0.004

Mencari setelan

Melatih variabel yang sudah dilatih dari SavedModel yang diimpor bersama dengan variabel model di sekitarnya disebut fine-tuning the SavedModel. Hal ini dapat menghasilkan kualitas yang lebih baik, tetapi seringkali membuat pelatihan lebih menuntut (mungkin membutuhkan lebih banyak waktu, lebih bergantung pada pengoptimal dan hyperparameternya, meningkatkan risiko overfitting dan memerlukan augmentasi kumpulan data, khususnya untuk CNN). Kami menyarankan konsumen SavedModel untuk melihat ke fine-tuning hanya setelah menetapkan rezim pelatihan yang baik, dan hanya jika penerbit SavedModel merekomendasikannya.

Penyesuaian halus mengubah parameter model "berkelanjutan" yang dilatih. Itu tidak mengubah transformasi hard-code, seperti tokenizing input teks dan pemetaan token ke entri yang sesuai dalam matriks embedding.

Untuk konsumen SavedModel

Membuat hub.KerasLayer seperti

layer = hub.KerasLayer(..., trainable=True)

memungkinkan fine-tuning dari SavedModel dimuat oleh lapisan. Ia menambahkan bobot yang dapat dilatih dan pengatur bobot yang dideklarasikan dalam SavedModel ke model Keras, dan menjalankan komputasi SavedModel dalam mode pelatihan (pikirkan tentang putus sekolah, dll.).

Colab klasifikasi gambar berisi contoh menyeluruh dengan penyempurnaan opsional.

Mengekspor ulang hasil fine-tuning

Pengguna tingkat lanjut mungkin ingin menyimpan hasil penyetelan ulang ke dalam SavedModel yang dapat digunakan sebagai pengganti yang dimuat sebelumnya. Ini dapat dilakukan dengan kode seperti

loaded_obj = hub.load("https://tfhub.dev/...")
hub_layer = hub.KerasLayer(loaded_obj, trainable=True, ...)

model = keras.Sequential([..., hub_layer, ...])
model.compile(...)
model.fit(...)

export_module_dir = os.path.join(os.getcwd(), "finetuned_model_export")
tf.saved_model.save(loaded_obj, export_module_dir)

Untuk pencipta SavedModel

Saat membuat SavedModel untuk dibagikan di TensorFlow Hub, pikirkan terlebih dahulu apakah dan bagaimana konsumen harus menyempurnakannya, dan berikan panduan dalam dokumentasi.

Menyimpan dari Model Keras harus membuat semua mekanisme fine-tuning bekerja (menghemat kerugian regularisasi bobot, mendeklarasikan variabel yang dapat dilatih, menelusuri __call__ untuk kedua training=True dan training=False , dll.)

Pilih antarmuka model yang berfungsi baik dengan aliran gradien, misalnya, log keluaran alih-alih probabilitas softmax atau prediksi top-k.

Jika model menggunakan dropout, normalisasi batch, atau teknik pelatihan serupa yang melibatkan hyperparameter, setel ke nilai yang masuk akal di banyak masalah target dan ukuran batch yang diharapkan. (Sampai tulisan ini dibuat, menabung dari Keras tidak memudahkan konsumen untuk menyesuaikannya.)

Regularizer berat pada setiap lapisan disimpan (dengan koefisien kekuatan regularisasi), tetapi regularisasi berat dari dalam pengoptimal (seperti tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) hilang. Beri tahu konsumen tentang SavedModel Anda.