TensorFlow 2 中來自 TF Hub 的 SavedModel

TensorFlow 2 的 SavedModel 格式是建議在 TensorFlow Hub 上分享預先訓練模型和模型組件的方式。它取代了舊版的 TF1 Hub 格式，並隨附一組新的 API。

本頁說明如何在 TensorFlow 2 程式中使用低階 hub.load() API 及其 hub.KerasLayer 包裝函式來重複使用 TF2 SavedModel。(通常，hub.KerasLayer 會與其他 tf.keras.layers 結合，以建構 Keras 模型或 TF2 Estimator 的 model_fn。) 這些 API 也可在限制範圍內載入 TF1 Hub 格式的舊版模型，詳情請參閱相容性指南。

TensorFlow 1 的使用者可以更新至 TF 1.15，然後使用相同的 API。舊版 TF1 無法運作。

使用來自 TF Hub 的 SavedModel

在 Keras 中使用 SavedModel

Keras 是 TensorFlow 的高階 API，可藉由組合 Keras Layer 物件來建構深度學習模型。tensorflow_hub 程式庫提供類別 hub.KerasLayer，此類別會使用 SavedModel 的網址 (或檔案系統路徑) 初始化，然後提供來自 SavedModel 的運算，包括其預先訓練的權重。

以下範例說明如何使用預先訓練的文字嵌入

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)

由此，可以透過一般 Keras 方式建構文字分類器

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

文字分類 Colab 完整範例說明如何訓練及評估這類分類器。

hub.KerasLayer 中的模型權重預設設為不可訓練。如要變更此設定，請參閱下方的微調章節。權重會在相同圖層物件的所有應用程式之間共用，這在 Keras 中很常見。

在 Estimator 中使用 SavedModel

TensorFlow Estimator API 的分散式訓練使用者可以透過 hub.KerasLayer 和其他 tf.keras.layers 撰寫 model_fn，藉此使用來自 TF Hub 的 SavedModel。

幕後花絮：SavedModel 下載與快取

使用來自 TensorFlow Hub (或其他實作其代管協定的 HTTPS 伺服器) 的 SavedModel 會將其下載並解壓縮到本機檔案系統 (如果尚未存在)。可以設定環境變數 TFHUB_CACHE_DIR，以覆寫快取已下載及未壓縮 SavedModel 的預設暫存位置。詳情請參閱快取。

在低階 TensorFlow 中使用 SavedModel

模型控制代碼

SavedModel 可以從指定的 handle 載入，其中 handle 是檔案系統路徑、有效的 TFhub.dev 模型網址 (例如「https://tfhub.dev/...」)。Kaggle 模型網址會根據我們的條款及與模型資產相關聯的授權，鏡像 TFhub.dev 控制代碼，例如「https://www.kaggle.com/...」。來自 Kaggle 模型的控制代碼與其對應的 TFhub.dev 控制代碼相同。

函式 hub.load(handle) 會下載並解壓縮 SavedModel (除非 handle 已是檔案系統路徑)，然後傳回使用 TensorFlow 內建函式 tf.saved_model.load() 載入的結果。因此，hub.load() 可以處理任何有效的 SavedModel (與 TF1 的前身 hub.Module 不同)。

進階主題：載入後對 SavedModel 的預期

根據 SavedModel 的內容，可以透過各種方式叫用 obj = hub.load(...) 的結果 (詳情請參閱 TensorFlow 的 SavedModel 指南)

SavedModel 的服務簽名 (如果有的話) 會表示為具體函式的字典，而且可以像 tensors_out = obj.signatures["serving_default"](**tensors_in) 一樣呼叫，其中張量的字典會以個別輸入和輸出名稱做為索引鍵，並受限於簽名的形狀和 dtype 限制。
已儲存物件的 @tf.function 裝飾方法 (如果有的話) 會還原為 tf.function 物件，這些物件可以針對 tf.function 在儲存前已追蹤的張量和非張量引數的所有組合呼叫。特別是，如果具有適用追蹤的 obj.__call__ 方法，則可以像 Python 函式一樣呼叫 obj 本身。簡單的範例可能如下所示：output_tensor = obj(input_tensor, training=False)。

這為 SavedModel 可以實作的介面留下極大的自由度。適用於 obj 的可重複使用 SavedModel 介面建立慣例，讓用戶端程式碼 (包括 hub.KerasLayer 等轉接器) 瞭解如何使用 SavedModel。

部分 SavedModel 可能未遵循該慣例，尤其是並非要在較大型模型中重複使用的完整模型，而僅提供服務簽名。

SavedModel 中的可訓練變數會重新載入為可訓練，且 tf.GradientTape 預設會監看這些變數。如需注意須知，請參閱下方的微調章節，並考慮避免初學者犯的錯誤。即使您想要微調，您可能也想查看 obj.trainable_variables 是否建議僅重新訓練原始可訓練變數的子集。

為 TF Hub 建立 SavedModel

總覽

SavedModel 是 TensorFlow 的標準序列化格式，適用於已訓練模型或模型組件。它會將模型的已訓練權重與執行其運算的確切 TensorFlow 運算儲存在一起。它可以獨立於建立它的程式碼使用。特別是，它可以跨不同的高階模型建構 API (例如 Keras) 重複使用，因為 TensorFlow 運算是它們通用的基本語言。

從 Keras 儲存

從 TensorFlow 2 開始，tf.keras.Model.save() 和 tf.keras.models.save_model() 預設為 SavedModel 格式 (而非 HDF5)。產生的 SavedModel 可以與 hub.load()、hub.KerasLayer 和類似的轉接器搭配使用，適用於其他高階 API (在這些 API 可用時)。

如要分享完整的 Keras 模型，只要使用 include_optimizer=False 儲存即可。

如要分享 Keras 模型的一部分，請將該部分設為模型本身，然後儲存。您可以從一開始就佈局程式碼....

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

...或在事後剪下要分享的部分 (如果它與完整模型的分層一致)

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)

GitHub 上的 TensorFlow 模型針對 BERT 使用前一種方法 (請參閱 nlp/tools/export_tfhub_lib.py，請注意 core_model (用於匯出) 與 pretrainer (用於還原檢查點) 之間的區隔)，而針對 ResNet 使用後一種方法 (請參閱 legacy/image_classification/tfhub_export.py)。

從低階 TensorFlow 儲存

這需要非常熟悉 TensorFlow 的 SavedModel 指南。

如果您想要提供的內容不只是服務簽名，則應實作可重複使用 SavedModel 介面。概念上，這看起來像

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

微調

將匯入的 SavedModel 的已訓練變數與周圍模型的變數一起訓練，稱為微調 SavedModel。這可以產生更好的品質，但通常會使訓練更耗費資源 (可能需要更多時間、更依賴最佳化工具及其超參數、增加過度擬合的風險，並需要資料集擴增，尤其是針對 CNN)。我們建議 SavedModel 消費者在建立良好的訓練機制後再研究微調，且僅在 SavedModel 發布者建議時才進行微調。

微調會變更已訓練的「連續」模型參數。它不會變更硬式編碼轉換，例如將文字輸入符號化，以及將符號對應至嵌入矩陣中的對應項目。

適用於 SavedModel 消費者

建立 hub.KerasLayer，例如

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

啟用圖層載入的 SavedModel 微調。它會將 SavedModel 中宣告的可訓練權重和權重正規化器新增至 Keras 模型，並在訓練模式中執行 SavedModel 的運算 (請考慮 dropout 等)。

圖片分類 Colab 包含具有選用微調功能的端對端範例。

重新匯出微調結果

進階使用者可能想要將微調結果儲存回 SavedModel，以取代原始載入的 SavedModel。這可以使用如下程式碼完成

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)

適用於 SavedModel 建立者

在建立 SavedModel 以在 TensorFlow Hub 上分享時，請預先思考其消費者是否應微調以及如何微調，並在文件中提供指引。

從 Keras 模型儲存應可讓微調的所有機制都能運作 (儲存權重正規化損失、宣告可訓練變數、追蹤 __call__ 以用於 training=True 和 training=False 等)

選擇與梯度流程搭配良好的模型介面，例如輸出 logits 而非 softmax 機率或前 k 項預測。

如果模型使用 dropout、批次正規化或類似涉及超參數的訓練技術，請將其設定為適用於許多預期目標問題和批次大小的值。(在撰寫本文時，從 Keras 儲存並不容易讓消費者調整它們。)

個別圖層上的權重正規化器會儲存 (及其正規化強度係數)，但來自最佳化工具內部的權重正規化 (例如 tf.keras.optimizers.Ftrl.l1_regularization_strength=...)) 會遺失。請相應地告知 SavedModel 的消費者。