SavedModels dari TF Hub di TensorFlow 2

Format SavedModel TensorFlow 2 adalah cara yang disarankan untuk berbagi model terlatih dan potongan model di TensorFlow Hub. Ini menggantikan format TF1 Hub yang lama dan dilengkapi dengan serangkaian API baru.

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

Pengguna TensorFlow 1 dapat mengupdate ke TF 1.15 lalu menggunakan API yang sama. TF1 versi 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 kelas 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 terlatih:

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 dibuat dengan cara Keras biasa:

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

Kolab klasifikasi Teks adalah contoh lengkap cara melatih dan mengevaluasi pengklasifikasi tersebut.

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

Menggunakan SavedModel di Estimator

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

Di balik layar: Pengunduhan dan cache SavedModel

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

Menggunakan SavedModel di TensorFlow tingkat rendah

Pegangan Model

SavedModels dapat dimuat dari handle tertentu, yang handle adalah jalur sistem file, URL model TFhub.dev yang valid (misalnya "https://tfhub.dev/..."). URL Model Kaggle mencerminkan pegangan TFhub.dev sesuai dengan Ketentuan kami dan lisensi yang terkait dengan aset model, misalnya, "https://www.kaggle.com/...". Pegangan dari Model Kaggle setara dengan pegangan TFhub.dev yang sesuai.

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

Topik lanjutan: apa yang diharapkan dari SavedModel setelah memuat

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

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

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

Hal ini memberikan kebebasan besar dalam antarmuka yang dapat diterapkan oleh SavedModels. Antarmuka SavedModels yang Dapat Digunakan Kembali untuk obj menetapkan konvensi sedemikian rupa sehingga kode klien, termasuk adaptor seperti hub.KerasLayer , mengetahui cara menggunakan SavedModel.

Beberapa SavedModel mungkin tidak mengikuti konvensi tersebut, 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 mengawasinya secara default. Lihat bagian penyesuaian di bawah untuk mengetahui beberapa peringatan, dan pertimbangkan untuk menghindari hal ini sebagai permulaan. Bahkan jika Anda ingin menyempurnakannya, Anda mungkin ingin melihat apakah obj.trainable_variables menyarankan untuk melatih ulang hanya sebagian dari variabel yang semula dapat dilatih.

Membuat SavedModels untuk TF Hub

Ringkasan

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

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 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 bagian dari Model Keras, buatlah bagian itu menjadi Model sendiri, lalu simpanlah. 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 kejadian (jika sejajar dengan lapisan 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

Hal ini memerlukan pemahaman yang baik tentang Panduan SavedModel TensorFlow.

Jika Anda ingin memberikan lebih dari sekedar 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

Penyempurnaan

Melatih variabel yang sudah dilatih dari SavedModel yang diimpor bersama dengan variabel model di sekitarnya disebut menyempurnakan SavedModel. Hal ini dapat menghasilkan kualitas yang lebih baik, namun sering kali membuat pelatihan lebih menuntut (mungkin memerlukan waktu lebih lama, lebih bergantung pada pengoptimal dan hyperparameternya, meningkatkan risiko overfitting, dan memerlukan augmentasi set data, khususnya untuk CNN). Kami menyarankan konsumen SavedModel untuk melakukan penyesuaian hanya setelah menetapkan rezim pelatihan yang baik, dan hanya jika penerbit SavedModel merekomendasikannya.

Penyempurnaan mengubah parameter model "kontinu" yang dilatih. Itu tidak mengubah transformasi hard-code, seperti memberi token pada input teks dan memetakan token ke entri yang sesuai dalam matriks penyematan.

Untuk konsumen SavedModel

Membuat hub.KerasLayer seperti

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

memungkinkan penyempurnaan SavedModel yang 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 dropout, dll.).

Kolab klasifikasi gambar berisi contoh menyeluruh dengan penyesuaian opsional.

Mengekspor ulang hasil fine tuning

Pengguna tingkat lanjut mungkin ingin menyimpan hasil penyesuaian kembali 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 pembuat 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 seharusnya membuat semua mekanisme penyesuaian berfungsi (menghemat penurunan regularisasi bobot, mendeklarasikan variabel yang dapat dilatih, menelusuri __call__ untuk training=True dan training=False , dll.)

Pilih antarmuka model yang cocok dengan aliran gradien, misalnya, log keluaran, bukan 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. (Saat tulisan ini dibuat, menabung dari Keras tidak memudahkan konsumen untuk menyesuaikannya.)

Pengatur bobot pada masing-masing lapisan disimpan (dengan koefisien kekuatan regularisasinya), tetapi regularisasi bobot dari dalam pengoptimal (seperti tf.keras.optimizers.Ftrl.l1_regularization_strength=...) ) hilang. Beri tahu konsumen tentang SavedModel Anda dengan tepat.