使用 SavedModel 格式

在 TensorFlow.org 上檢視

在 Google Colab 中執行

在 GitHub 上檢視原始碼

下載筆記本

SavedModel 包含完整的 TensorFlow 程式，包括已訓練的參數 (即 tf.Variable) 和運算。它不需要執行原始模型建構程式碼，因此適用於與 TFLite、TensorFlow.js、TensorFlow Serving 或 TensorFlow Hub 共用或部署。

您可以使用下列 API，以 SavedModel 格式儲存及載入模型

• 低階 tf.saved_model API。本文件詳細說明如何使用此 API。
  • 儲存：tf.saved_model.save(model, path_to_dir)
  • 載入：model = tf.saved_model.load(path_to_dir)
• 高階 tf.keras.Model API。請參閱Keras 儲存與序列化指南。
• 如果您只想在訓練期間儲存/載入權重，請參閱檢查點指南。

從 Keras 建立 SavedModel

為了快速入門，本節將匯出預先訓練的 Keras 模型，並使用它提供圖片分類要求。本指南的其餘部分將詳細說明並討論建立 SavedModel 的其他方式。

import os
import tempfile

from matplotlib import pyplot as plt
import numpy as np
import tensorflow as tf

tmpdir = tempfile.mkdtemp()

physical_devices = tf.config.list_physical_devices('GPU')
for device in physical_devices:
  tf.config.experimental.set_memory_growth(device, True)

file = tf.keras.utils.get_file(
    "grace_hopper.jpg",
    "https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg")
img = tf.keras.utils.load_img(file, target_size=[224, 224])
plt.imshow(img)
plt.axis('off')
x = tf.keras.utils.img_to_array(img)
x = tf.keras.applications.mobilenet.preprocess_input(
    x[tf.newaxis,...])

您將使用葛麗絲·霍普上將的圖片作為執行範例，以及 Keras 預先訓練的圖片分類模型，因為它很容易使用。自訂模型也適用，稍後會詳細介紹。

labels_path = tf.keras.utils.get_file(
    'ImageNetLabels.txt',
    'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())

pretrained_model = tf.keras.applications.MobileNet()
result_before_save = pretrained_model(x)

decoded = imagenet_labels[np.argsort(result_before_save)[0,::-1][:5]+1]

print("Result before saving:\n", decoded)

此圖片的最佳預測是「軍裝」。

mobilenet_save_path = os.path.join(tmpdir, "mobilenet/1/")
tf.saved_model.save(pretrained_model, mobilenet_save_path)

儲存路徑遵循 TensorFlow Serving 使用的慣例，其中最後一個路徑元件 (此處的 1/) 是模型的版本號碼 - 它允許 Tensorflow Serving 等工具推論相對的新鮮度。

您可以使用 tf.saved_model.load 將 SavedModel 載回 Python，並查看如何分類霍普上將的圖片。

loaded = tf.saved_model.load(mobilenet_save_path)
print(list(loaded.signatures.keys()))  # ["serving_default"]

匯入的簽名一律會傳回字典。如要自訂簽名名稱和輸出字典鍵，請參閱在匯出期間指定簽名。

infer = loaded.signatures["serving_default"]
print(infer.structured_outputs)

從 SavedModel 執行推論會產生與原始模型相同的結果。

labeling = infer(tf.constant(x))[pretrained_model.output_names[0]]

decoded = imagenet_labels[np.argsort(labeling)[0,::-1][:5]+1]

print("Result after saving and loading:\n", decoded)

在 TensorFlow Serving 中執行 SavedModel

SavedModel 可從 Python 使用 (稍後會詳細說明)，但生產環境通常會使用專用服務進行推論，而無需執行 Python 程式碼。這很容易使用 TensorFlow Serving 從 SavedModel 進行設定。

如需端對端 tensorflow-serving 範例，請參閱TensorFlow Serving REST 教學課程。

磁碟上的 SavedModel 格式

SavedModel 是一個目錄，其中包含序列化的簽名和執行這些簽名所需的狀態，包括變數值和詞彙表。

ls {mobilenet_save_path}

saved_model.pb 檔案儲存實際的 TensorFlow 程式或模型，以及一組具名的簽名，每個簽名都識別接受張量輸入並產生張量輸出的函式。

SavedModel 可能包含模型的多個變體 (多個 v1.MetaGraphDefs，以 --tag_set 標記至 saved_model_cli 識別)，但這種情況很少見。建立模型多個變體的 API 包括 tf.Estimator.experimental_export_all_saved_models 和 TensorFlow 1.x 中的 tf.saved_model.Builder。

saved_model_cli show --dir {mobilenet_save_path} --tag_set serve

variables 目錄包含標準訓練檢查點 (請參閱訓練檢查點指南)。

ls {mobilenet_save_path}/variables

assets 目錄包含 TensorFlow 圖使用的檔案，例如用於初始化詞彙表表格的文字檔。在此範例中未使用。

SavedModel 可能具有 assets.extra 目錄，用於 TensorFlow 圖未使用的任何檔案，例如關於如何處理 SavedModel 的消費者資訊。TensorFlow 本身不使用此目錄。

fingerprint.pb 檔案包含 SavedModel 的指紋，它由多個 64 位元雜湊組成，可唯一識別 SavedModel 的內容。指紋 API 目前為實驗性，但 tf.saved_model.experimental.read_fingerprint 可用於將 SavedModel 指紋讀取到 tf.saved_model.experimental.Fingerprint 物件中。

儲存自訂模型

tf.saved_model.save 支援儲存 tf.Module 物件及其子類別，例如 tf.keras.Layer 和 tf.keras.Model。

讓我們看看儲存和還原 tf.Module 的範例。

class CustomModule(tf.Module):

  def __init__(self):
    super(CustomModule, self).__init__()
    self.v = tf.Variable(1.)

  @tf.function
  def __call__(self, x):
    print('Tracing with', x)
    return x * self.v

  @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
  def mutate(self, new_v):
    self.v.assign(new_v)

module = CustomModule()

當您儲存 tf.Module 時，會儲存任何 tf.Variable 屬性、tf.function 裝飾的方法，以及透過遞迴遍歷找到的 tf.Module。(如要進一步瞭解此遞迴遍歷，請參閱檢查點教學課程。) 但是，任何 Python 屬性、函式和資料都會遺失。這表示當 tf.function 儲存時，不會儲存任何 Python 程式碼。

如果未儲存任何 Python 程式碼，SavedModel 如何知道如何還原函式？

簡而言之，tf.function 的運作方式是追蹤 Python 程式碼以產生 ConcreteFunction (圍繞 tf.Graph 的可呼叫包裝函式)。當儲存 tf.function 時，您實際上是在儲存 tf.function 的 ConcreteFunction 快取。

如要進一步瞭解 tf.function 和 ConcreteFunction 之間的關係，請參閱 tf.function 指南。

module_no_signatures_path = os.path.join(tmpdir, 'module_no_signatures')
module(tf.constant(0.))
print('Saving model...')
tf.saved_model.save(module, module_no_signatures_path)

載入和使用自訂模型

當您在 Python 中載入 SavedModel 時，所有 tf.Variable 屬性、tf.function 裝飾的方法和 tf.Module 都會以與原始儲存的 tf.Module 相同的物件結構還原。

imported = tf.saved_model.load(module_no_signatures_path)
assert imported(tf.constant(3.)).numpy() == 3
imported.mutate(tf.constant(2.))
assert imported(tf.constant(3.)).numpy() == 6

因為未儲存任何 Python 程式碼，所以使用新的輸入簽名呼叫 tf.function 將會失敗

imported(tf.constant([3.]))

ValueError: Could not find matching function to call for canonicalized inputs ((,), {}). Only existing signatures are [((TensorSpec(shape=(), dtype=tf.float32, name=u'x'),), {})].

基本微調

變數物件可用，您可以透過匯入的函式進行反向傳播。這足以在簡單情況下微調 (即重新訓練) SavedModel。

optimizer = tf.keras.optimizers.SGD(0.05)

def train_step():
  with tf.GradientTape() as tape:
    loss = (10. - imported(tf.constant(2.))) ** 2
  variables = tape.watched_variables()
  grads = tape.gradient(loss, variables)
  optimizer.apply_gradients(zip(grads, variables))
  return loss

for _ in range(10):
  # "v" approaches 5, "loss" approaches 0
  print("loss={:.2f} v={:.2f}".format(train_step(), imported.v.numpy()))

一般微調

來自 Keras 的 SavedModel 提供比純 __call__ 更多詳細資訊，以解決更進階的微調案例。TensorFlow Hub 建議在為微調目的共用的 SavedModel 中提供下列項目 (如果適用)

• 如果模型使用 Dropout 或其他技術，其中正向傳遞在訓練和推論之間有所不同 (例如批次正規化)，則 __call__ 方法會採用選用的 Python 值 training= 引數，預設為 False，但可以設定為 True。
• 在 __call__ 屬性旁邊，還有 .variable 和 .trainable_variable 屬性，其中包含對應的變數清單。最初可訓練但打算在微調期間凍結的變數會從 .trainable_variables 中省略。
• 為了像 Keras 這樣的架構 (將權重正規化器表示為層或子模型的屬性)，也可以有 .regularization_losses 屬性。它包含零引數函式的清單，其值旨在新增至總損失。

回到最初的 MobileNet 範例，您可以看到其中一些在運作

loaded = tf.saved_model.load(mobilenet_save_path)
print("MobileNet has {} trainable variables: {}, ...".format(
          len(loaded.trainable_variables),
          ", ".join([v.name for v in loaded.trainable_variables[:5]])))

trainable_variable_ids = {id(v) for v in loaded.trainable_variables}
non_trainable_variables = [v for v in loaded.variables
                           if id(v) not in trainable_variable_ids]
print("MobileNet also has {} non-trainable variables: {}, ...".format(
          len(non_trainable_variables),
          ", ".join([v.name for v in non_trainable_variables[:3]])))

在匯出期間指定簽名

TensorFlow Serving 和 saved_model_cli 等工具可以與 SavedModel 互動。為了協助這些工具判斷要使用哪些 ConcreteFunction，您需要指定服務簽名。tf.keras.Model 會自動指定服務簽名，但您必須為我們的自訂模組明確宣告服務簽名。

根據預設，自訂 tf.Module 中未宣告任何簽名。

assert len(imported.signatures) == 0

如要宣告服務簽名，請使用 signatures kwarg 指定 ConcreteFunction。當指定單一簽名時，其簽名金鑰將為 'serving_default'，它會儲存為常數 tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY。

module_with_signature_path = os.path.join(tmpdir, 'module_with_signature')
call = module.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
tf.saved_model.save(module, module_with_signature_path, signatures=call)

imported_with_signatures = tf.saved_model.load(module_with_signature_path)
list(imported_with_signatures.signatures.keys())

如要匯出多個簽名，請將簽名金鑰字典傳遞至 ConcreteFunction。每個簽名金鑰都對應一個 ConcreteFunction。

module_multiple_signatures_path = os.path.join(tmpdir, 'module_with_multiple_signatures')
signatures = {"serving_default": call,
              "array_input": module.__call__.get_concrete_function(tf.TensorSpec([None], tf.float32))}

tf.saved_model.save(module, module_multiple_signatures_path, signatures=signatures)

imported_with_multiple_signatures = tf.saved_model.load(module_multiple_signatures_path)
list(imported_with_multiple_signatures.signatures.keys())

根據預設，輸出張量名稱相當通用，例如 output_0。如要控制輸出的名稱，請修改您的 tf.function 以傳回將輸出名稱對應至輸出的字典。輸入的名稱衍生自 Python 函式引數名稱。

class CustomModuleWithOutputName(tf.Module):
  def __init__(self):
    super(CustomModuleWithOutputName, self).__init__()
    self.v = tf.Variable(1.)

  @tf.function(input_signature=[tf.TensorSpec(None, tf.float32)])
  def __call__(self, x):
    return {'custom_output_name': x * self.v}

module_output = CustomModuleWithOutputName()
call_output = module_output.__call__.get_concrete_function(tf.TensorSpec(None, tf.float32))
module_output_path = os.path.join(tmpdir, 'module_with_output_name')
tf.saved_model.save(module_output, module_output_path,
                    signatures={'serving_default': call_output})

imported_with_output_name = tf.saved_model.load(module_output_path)
imported_with_output_name.signatures['serving_default'].structured_outputs

Proto 分割

由於 protobuf 實作的限制，proto 大小不能超過 2GB。當嘗試儲存非常大的模型時，可能會導致下列錯誤

ValueError: Message tensorflow.SavedModel exceeds maximum protobuf size of 2GB: ...

google.protobuf.message.DecodeError: Error parsing message as the message exceeded the protobuf limit with type 'tensorflow.GraphDef'

如果您想要儲存超過 2GB 限制的模型，則需要使用新的 proto 分割選項儲存

tf.saved_model.save(
  ...,
  options=tf.saved_model.SaveOptions(experimental_image_format=True)
)

如需更多資訊，請參閱Proto Splitter / Merger 程式庫指南。

在 C++ 中載入 SavedModel

C++ 版本的 SavedModel 載入器提供 API 從路徑載入 SavedModel，同時允許 SessionOptions 和 RunOptions。您必須指定與要載入的圖相關聯的標記。載入版本的 SavedModel 稱為 SavedModelBundle，其中包含 MetaGraphDef 和已載入它的工作階段。

const string export_dir = ...
SavedModelBundle bundle;
...
LoadSavedModel(session_options, run_options, export_dir, {kSavedModelTagTrain},
               &bundle);

$ bazel build //tensorflow/python/tools:saved_model_cli

SavedModel 包含一或多個模型變體 (技術上來說，v1.MetaGraphDef)，由其標記集識別。為了服務模型，您可能會想知道每個模型變體中有哪些 SignatureDef，以及它們的輸入和輸出是什麼。show 命令可讓您以階層式順序檢查 SavedModel 的內容。以下是語法

usage: saved_model_cli show [-h] --dir DIR [--all]
[--tag_set TAG_SET] [--signature_def SIGNATURE_DEF_KEY]

例如，下列命令顯示 SavedModel 中所有可用的標記集

$ saved_model_cli show --dir /tmp/saved_model_dir
The given SavedModel contains the following tag-sets:
serve
serve, gpu

下列命令顯示標記集的所有可用 SignatureDef 金鑰

$ saved_model_cli show --dir /tmp/saved_model_dir --tag_set serve
The given SavedModel `MetaGraphDef` contains `SignatureDefs` with the
following keys:
SignatureDef key: "classify_x2_to_y3"
SignatureDef key: "classify_x_to_y"
SignatureDef key: "regress_x2_to_y3"
SignatureDef key: "regress_x_to_y"
SignatureDef key: "regress_x_to_y2"
SignatureDef key: "serving_default"

如果標記集中有多個標記，則必須指定所有標記，每個標記以逗號分隔。例如

$ saved_model_cli show --dir /tmp/saved_model_dir --tag_set serve,gpu

如要顯示特定 SignatureDef 的所有輸入和輸出 TensorInfo，請將 SignatureDef 金鑰傳遞至 signature_def 選項。當您想要知道張量金鑰值、dtype 和輸入張量的形狀，以便稍後執行運算圖時，這非常有用。例如

$ saved_model_cli show --dir \
/tmp/saved_model_dir --tag_set serve --signature_def serving_default
The given SavedModel SignatureDef contains the following input(s):
  inputs['x'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 1)
      name: x:0
The given SavedModel SignatureDef contains the following output(s):
  outputs['y'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 1)
      name: y:0
Method name is: tensorflow/serving/predict

如要顯示 SavedModel 中的所有可用資訊，請使用 --all 選項。例如

$ saved_model_cli show --dir /tmp/saved_model_dir --all
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['classify_x2_to_y3']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['inputs'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: x2:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['scores'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: y3:0
  Method name is: tensorflow/serving/classify

...

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['x'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: x:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['y'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: y:0
  Method name is: tensorflow/serving/predict

run 命令

叫用 run 命令以執行圖運算、傳遞輸入，然後顯示 (並選擇性地儲存) 輸出。以下是語法

usage: saved_model_cli run [-h] --dir DIR --tag_set TAG_SET --signature_def
                           SIGNATURE_DEF_KEY [--inputs INPUTS]
                           [--input_exprs INPUT_EXPRS]
                           [--input_examples INPUT_EXAMPLES] [--outdir OUTDIR]
                           [--overwrite] [--tf_debug]

run 命令提供下列三種方式來將輸入傳遞至模型

• --inputs 選項可讓您在檔案中傳遞 numpy ndarray。
• --input_exprs 選項可讓您傳遞 Python 運算式。
• --input_examples 選項可讓您傳遞 tf.train.Example。

--inputs

如要在檔案中傳遞輸入資料，請指定 --inputs 選項，它採用下列一般格式

--inputs <INPUTS>

其中 INPUTS 是下列其中一種格式

• <input_key>=<filename>
• <input_key>=<filename>[<variable_name>]

您可以傳遞多個 INPUTS。如果您確實傳遞多個輸入，請使用分號分隔每個 INPUTS。

saved_model_cli 使用 numpy.load 來載入 filename。filename 可以是下列任何格式

• .npy
• .npz
• pickle 格式

.npy 檔案一律包含 numpy ndarray。因此，從 .npy 檔案載入時，內容將直接指派給指定的輸入張量。如果您使用該 .npy 檔案指定 variable_name，則會忽略 variable_name 並發出警告。

從 .npz (zip) 檔案載入時，您可以選擇性地指定 variable_name，以識別 zip 檔案中要為輸入張量金鑰載入的變數。如果您未指定 variable_name，SavedModel CLI 將檢查 zip 檔案中是否僅包含一個檔案，並為指定的輸入張量金鑰載入它。

從 pickle 檔案載入時，如果在方括號中未指定 variable_name，則 pickle 檔案中的任何內容都將傳遞至指定的輸入張量金鑰。否則，SavedModel CLI 會假設字典儲存在 pickle 檔案中，並使用對應於 variable_name 的值。

--input_exprs

如要透過 Python 運算式傳遞輸入，請指定 --input_exprs 選項。當您沒有周圍的資料檔案，但仍想使用一些符合模型 SignatureDef 的 dtype 和形狀的簡單輸入來執行模型健全性檢查時，這會很有用。例如

`<input_key>=[[1],[2],[3]]`

除了 Python 運算式之外，您也可以傳遞 numpy 函式。例如

`<input_key>=np.ones((32,32,3))`

(請注意，numpy 模組已以 np 的形式提供給您。)

--input_examples

如要傳遞 tf.train.Example 作為輸入，請指定 --input_examples 選項。對於每個輸入金鑰，它會採用字典清單，其中每個字典都是 tf.train.Example 的執行個體。字典金鑰是特徵，值是每個特徵的值清單。例如

`<input_key>=[{"age":[22,24],"education":["BS","MS"]}]`

儲存輸出

根據預設，SavedModel CLI 會將輸出寫入標準輸出。如果目錄傳遞至 --outdir 選項，則輸出將儲存為 .npy 檔案，並以給定目錄下的輸出張量金鑰命名。

使用 --overwrite 覆寫現有的輸出檔案。