遷移 SavedModel 工作流程

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

將模型從 TensorFlow 1 的圖表和工作階段遷移至 TensorFlow 2 API (例如 tf.function、tf.Module 和 tf.keras.Model) 後，即可遷移模型儲存和載入程式碼。本筆記本提供如何在 TensorFlow 1 和 TensorFlow 2 中以 SavedModel 格式儲存和載入的範例。以下快速總覽從 TensorFlow 1 遷移至 TensorFlow 2 的相關 API 變更

TensorFlow 1 遷移至 TensorFlow 2
儲存 tf.compat.v1.saved_model.Builder
tf.compat.v1.saved_model.simple_save		 tf.saved_model.save
Keras：tf.keras.models.save_model
載入 tf.compat.v1.saved_model.load tf.saved_model.load
Keras：tf.keras.models.load_model
簽名：一組輸入
和輸出張量，
可用於執行
 使用 *.signature_def 公用程式產生
(例如 tf.compat.v1.saved_model.predict_signature_def)		 寫入 tf.function 並使用 signatures 引數匯出
在 tf.saved_model.save 中。
分類
和迴歸:
特殊類型的簽名		 使用下列項目產生
tf.compat.v1.saved_model.classification_signature_def,
tf.compat.v1.saved_model.regression_signature_def,
和特定 Estimator 匯出。		 這兩種簽名類型已從 TensorFlow 2 中移除。
如果服務程式庫需要這些方法名稱，
tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater.

如需更深入的對應說明，請參閱下方的「從 TensorFlow 1 到 TensorFlow 2 的變更」一節。

設定

以下範例說明如何使用 TensorFlow 1 和 TensorFlow 2 API 將相同的虛擬 TensorFlow 模型 (在下方定義為 add_two) 匯出和載入為 SavedModel 格式。首先設定匯入項目和公用程式函式

import tensorflow as tf
import tensorflow.compat.v1 as tf1
import shutil

def remove_dir(path):
  try:
    shutil.rmtree(path)
  except:
    pass

def add_two(input):
  return input + 2

TensorFlow 1：儲存並匯出 SavedModel

在 TensorFlow 1 中，您可以使用 tf.compat.v1.saved_model.Builder、tf.compat.v1.saved_model.simple_save 和 tf.estimator.Estimator.export_saved_model API 來建構、儲存和匯出 TensorFlow 圖表和工作階段

1. 使用 SavedModelBuilder 將圖表儲存為 SavedModel

remove_dir("saved-model-builder")

with tf.Graph().as_default() as g:
  with tf1.Session() as sess:
    input = tf1.placeholder(tf.float32, shape=[])
    output = add_two(input)
    print("add two output: ", sess.run(output, {input: 3.}))

    # Save with SavedModelBuilder
    builder = tf1.saved_model.Builder('saved-model-builder')
    sig_def = tf1.saved_model.predict_signature_def(
        inputs={'input': input},
        outputs={'output': output})
    builder.add_meta_graph_and_variables(
        sess, tags=["serve"], signature_def_map={
            tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: sig_def
    })
    builder.save()

!saved_model_cli run --dir saved-model-builder --tag_set serve \
 --signature_def serving_default --input_exprs input=10

2. 建構用於服務的 SavedModel

remove_dir("simple-save")

with tf.Graph().as_default() as g:
  with tf1.Session() as sess:
    input = tf1.placeholder(tf.float32, shape=[])
    output = add_two(input)
    print("add_two output: ", sess.run(output, {input: 3.}))

    tf1.saved_model.simple_save(
        sess, 'simple-save',
        inputs={'input': input},
        outputs={'output': output})

!saved_model_cli run --dir simple-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10

3. 將 Estimator 推論圖表匯出為 SavedModel

在 Estimator model_fn (在下方定義) 的定義中，您可以藉由在 tf.estimator.EstimatorSpec 中傳回 export_outputs，在模型中定義簽名。輸出有不同的類型

  • tf.estimator.export.ClassificationOutput
  • tf.estimator.export.RegressionOutput
  • tf.estimator.export.PredictOutput

這些將分別產生分類、迴歸和預測簽名類型。

當 Estimator 使用 tf.estimator.Estimator.export_saved_model 匯出時，這些簽名將與模型一起儲存。

def model_fn(features, labels, mode):
  output = add_two(features['input'])
  step = tf1.train.get_global_step()
  return tf.estimator.EstimatorSpec(
      mode,
      predictions=output,
      train_op=step.assign_add(1),
      loss=tf.constant(0.),
      export_outputs={
          tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: \
          tf.estimator.export.PredictOutput({'output': output})})
est = tf.estimator.Estimator(model_fn, 'estimator-checkpoints')

# Train for one step to create a checkpoint.
def train_fn():
  return tf.data.Dataset.from_tensors({'input': 3.})
est.train(train_fn, steps=1)

# This utility function `build_raw_serving_input_receiver_fn` takes in raw
# tensor features and builds an "input serving receiver function", which
# creates placeholder inputs to the model.
serving_input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn(
    {'input': tf.constant(3.)})  # Pass in a dummy input batch.
estimator_path = est.export_saved_model('exported-estimator', serving_input_fn)

# Estimator's export_saved_model creates a time stamped directory. Move this
# to a set path so it can be inspected with `saved_model_cli` in the cell below.
!rm -rf estimator-model
import shutil
shutil.move(estimator_path, 'estimator-model')

!saved_model_cli run --dir estimator-model --tag_set serve \
 --signature_def serving_default --input_exprs input=[10]

TensorFlow 2：儲存並匯出 SavedModel

儲存並匯出以 tf.Module 定義的 SavedModel

若要在 TensorFlow 2 中匯出模型，您必須定義 tf.Module 或 tf.keras.Model，以保存模型的所有變數和函式。接著，您可以呼叫 tf.saved_model.save 來建立 SavedModel。如需瞭解詳情，請參閱「使用 SavedModel 格式」指南中的「儲存自訂模型」一節。

class MyModel(tf.Module):
  @tf.function
  def __call__(self, input):
    return add_two(input)

model = MyModel()

@tf.function
def serving_default(input):
  return {'output': model(input)}

signature_function = serving_default.get_concrete_function(
    tf.TensorSpec(shape=[], dtype=tf.float32))
tf.saved_model.save(
    model, 'tf2-save', signatures={
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_function})

!saved_model_cli run --dir tf2-save --tag_set serve \
 --signature_def serving_default --input_exprs input=10

儲存並匯出以 Keras 定義的 SavedModel

用於儲存和匯出的 Keras API (—Model.save 或 tf.keras.models.save_model—) 可以從 tf.keras.Model 匯出 SavedModel。如需更多詳細資訊，請查看「儲存和載入 Keras 模型」。

inp = tf.keras.Input(3)
out = add_two(inp)
model = tf.keras.Model(inputs=inp, outputs=out)

@tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.float32)])
def serving_default(input):
  return {'output': model(input)}

model.save('keras-model', save_format='tf', signatures={
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY: serving_default})

!saved_model_cli run --dir keras-model --tag_set serve \
 --signature_def serving_default --input_exprs input=10

載入 SavedModel

使用上述任何 API 儲存的 SavedModel，都可以使用 TensorFlow 1 或 TensorFlow 2 API 載入。

TensorFlow 1 SavedModel 通常可用於在載入 TensorFlow 2 時進行推論，但只有當 SavedModel 包含資源變數時，才能進行訓練 (產生梯度)。您可以檢查變數的 dtype；如果變數 dtype 包含「_ref」，則表示它是參考變數。

TensorFlow 2 SavedModel 可以從 TensorFlow 1 載入和執行，前提是 SavedModel 是使用簽名儲存的。

以下章節包含程式碼範例，說明如何載入先前章節中儲存的 SavedModel，並呼叫匯出的簽名。

TensorFlow 1：使用 tf.saved_model.load 載入 SavedModel

在 TensorFlow 1 中，您可以使用 tf.saved_model.load 將 SavedModel 直接匯入至目前的圖表和工作階段。您可以對張量輸入和輸出名稱呼叫 Session.run

def load_tf1(path, input):
  print('Loading from', path)
  with tf.Graph().as_default() as g:
    with tf1.Session() as sess:
      meta_graph = tf1.saved_model.load(sess, ["serve"], path)
      sig_def = meta_graph.signature_def[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
      input_name = sig_def.inputs['input'].name
      output_name = sig_def.outputs['output'].name
      print('  Output with input', input, ': ', 
            sess.run(output_name, feed_dict={input_name: input}))

load_tf1('saved-model-builder', 5.)
load_tf1('simple-save', 5.)
load_tf1('estimator-model', [5.])  # Estimator's input must be batched.
load_tf1('tf2-save', 5.)
load_tf1('keras-model', 5.)

TensorFlow 2：載入使用 tf.saved_model 儲存的模型

在 TensorFlow 2 中，物件會載入至 Python 物件，以儲存變數和函式。這與從 TensorFlow 1 儲存的模型相容。

如需瞭解詳情，請查看 tf.saved_model.load API 文件和 「載入和使用自訂模型」一節，此節取自「使用 SavedModel 格式」指南。

def load_tf2(path, input):
  print('Loading from', path)
  loaded = tf.saved_model.load(path)
  out = loaded.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY](
      tf.constant(input))['output']
  print('  Output with input', input, ': ', out)

load_tf2('saved-model-builder', 5.)
load_tf2('simple-save', 5.)
load_tf2('estimator-model', [5.])  # Estimator's input must be batched.
load_tf2('tf2-save', 5.)
load_tf2('keras-model', 5.)

使用 TensorFlow 2 API 儲存的模型也可以存取 tf.function 和附加至模型的變數 (而非匯出為簽名的變數)。例如

loaded = tf.saved_model.load('tf2-save')
print('restored __call__:', loaded.__call__)
print('output with input 5.', loaded(5))

TensorFlow 2：載入使用 Keras 儲存的模型

Keras 載入 API (—tf.keras.models.load_model—) 可讓您將儲存的模型重新載入回 Keras 模型物件。請注意，這僅允許您載入使用 Keras 儲存的 SavedModel (Model.save 或 tf.keras.models.save_model)。

使用 tf.saved_model.save 儲存的模型應使用 tf.saved_model.load 載入。您可以使用 tf.saved_model.load 載入使用 Model.save 儲存的 Keras 模型，但您只會取得 TensorFlow 圖表。如需瞭解詳情，請參閱 tf.keras.models.load_model API 文件和 「儲存和載入 Keras 模型」指南。

loaded_model = tf.keras.models.load_model('keras-model')
loaded_model.predict_on_batch(tf.constant([1, 3, 4]))

GraphDef 和 MetaGraphDef

沒有直接的方式可將原始 GraphDef 或 MetaGraphDef 載入至 TF2。不過，您可以使用 v1.wrap_function，將匯入圖表的 TF1 程式碼轉換為 TF2 concrete_function。

首先，儲存 MetaGraphDef

# Save a simple multiplication computation:
with tf.Graph().as_default() as g:
  x = tf1.placeholder(tf.float32, shape=[], name='x')
  v = tf.Variable(3.0, name='v')
  y = tf.multiply(x, v, name='y')
  with tf1.Session() as sess:
    sess.run(v.initializer)
    print(sess.run(y, feed_dict={x: 5}))
    s = tf1.train.Saver()
    s.export_meta_graph('multiply.pb', as_text=True)
    s.save(sess, 'multiply_values.ckpt')

使用 TF1 API，您可以使用 tf1.train.import_meta_graph 匯入圖表並還原值

with tf.Graph().as_default() as g:
  meta = tf1.train.import_meta_graph('multiply.pb')
  x = g.get_tensor_by_name('x:0')
  y = g.get_tensor_by_name('y:0')
  with tf1.Session() as sess:
    meta.restore(sess, 'multiply_values.ckpt')
    print(sess.run(y, feed_dict={x: 5}))

沒有用於載入圖表的 TF2 API，但您仍然可以將其匯入至可在即時模式中執行的具體函式

def import_multiply():
  # Any graph-building code is allowed here.
  tf1.train.import_meta_graph('multiply.pb')

# Creates a tf.function with all the imported elements in the function graph.
wrapped_import = tf1.wrap_function(import_multiply, [])
import_graph = wrapped_import.graph
x = import_graph.get_tensor_by_name('x:0')
y = import_graph.get_tensor_by_name('y:0')

# Restore the variable values.
tf1.train.Saver(wrapped_import.variables).restore(
    sess=None, save_path='multiply_values.ckpt')

# Create a concrete function by pruning the wrap_function (similar to sess.run).
multiply_fn = wrapped_import.prune(feeds=x, fetches=y)

# Run this function
multiply_fn(tf.constant(5.))  # inputs to concrete functions must be Tensors.

從 TensorFlow 1 到 TensorFlow 2 的變更

本節列出 TensorFlow 1 的主要儲存和載入詞彙、對等的 TensorFlow 2 詞彙，以及變更內容。

SavedModel

SavedModel 是一種格式，可儲存包含參數和運算的完整 TensorFlow 程式。其中包含服務平台用來執行模型的簽名。

檔案格式本身沒有顯著變更，因此 SavedModel 可以使用 TensorFlow 1 或 TensorFlow 2 API 載入和服務。

TensorFlow 1 和 TensorFlow 2 之間的差異

除了 API 變更之外，TensorFlow 2 中未更新 服務和推論使用案例，改進之處在於能夠重複使用和組合從 SavedModel 載入的模型。

在 TensorFlow 2 中，程式由物件 (例如 tf.Variable、tf.Module) 或更高階的 Keras 模型 (tf.keras.Model) 和層 (tf.keras.layers) 表示。不再有全域變數將值儲存在工作階段中，而且圖表現在存在於不同的 tf.function 中。因此，在模型匯出期間，SavedModel 會個別儲存每個元件和函式圖表。

當您使用 TensorFlow Python API 撰寫 TensorFlow 程式時，必須建構物件來管理變數、函式和其他資源。一般而言，這可透過使用 Keras API 來完成，但您也可以藉由建立或子類別化 tf.Module 來建構物件。

Keras 模型 (tf.keras.Model) 和 tf.Module 會自動追蹤附加至其中的變數和函式。SavedModel 會儲存模組、變數和函式之間的這些連線，以便在載入時還原這些連線。

簽名

簽名是 SavedModel 的端點，它們會告訴使用者如何執行模型以及需要哪些輸入。

在 TensorFlow 1 中，簽名是藉由列出輸入和輸出張量來建立。在 TensorFlow 2 中，簽名是藉由傳入具體函式來產生。(如需深入瞭解 TensorFlow 函式，請參閱「圖表和 tf.function 簡介」指南，特別是「多型性：一個函式，多個圖表」一節。) 簡而言之，具體函式是從 tf.function 產生。

# Option 1: Specify an input signature.
@tf.function(input_signature=[...])
def fn(...):
  ...
  return outputs

tf.saved_model.save(model, path, signatures={
    'name': fn
})

# Option 2: Call `get_concrete_function`
@tf.function
def fn(...):
  ...
  return outputs

tf.saved_model.save(model, path, signatures={
    'name': fn.get_concrete_function(...)
})

Session.run

在 TensorFlow 1 中，只要您已經知道張量名稱，就可以使用匯入的圖表呼叫 Session.run。這可讓您擷取還原的變數值，或執行模型中未在簽名中匯出的部分。

在 TensorFlow 2 中，您可以直接存取變數，例如權重矩陣 (kernel)

model = tf.Module()
model.dense_layer = tf.keras.layers.Dense(...)
tf.saved_model.save('my_saved_model')
loaded = tf.saved_model.load('my_saved_model')
loaded.dense_layer.kernel

或呼叫附加至模型物件的 tf.function：例如 loaded.__call__。

與 TF1 不同，沒有方法可以擷取函式的一部分並存取中繼值。您必須在儲存的物件中匯出所有需要的功能。

TensorFlow Serving 遷移注意事項

SavedModel 最初是為了與 TensorFlow Serving 搭配使用而建立。此平台提供不同類型的預測要求：分類、迴歸和預測。

TensorFlow 1 API 可讓您使用公用程式建立這些類型的簽名

分類 (classification_signature_def) 和 迴歸 (regression_signature_def) 限制了輸入和輸出，因此輸入必須是 tf.Example，而輸出必須是 classes、scores 或 prediction。同時，預測簽名 (predict_signature_def) 沒有任何限制。

使用 TensorFlow 2 API 匯出的 SavedModel 與 TensorFlow Serving 相容，但只會包含預測簽名。分類和迴歸簽名已移除。

如果您需要使用分類和迴歸簽名，可以使用 tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater 修改匯出的 SavedModel。

後續步驟

如需進一步瞭解 TensorFlow 2 中的 SavedModel，請查看下列指南

如果您使用 TensorFlow Hub，可能會發現這些指南很有用