Python tf.saved_model.load用法及代碼示例

從 export_dir 加載 SavedModel。

用法

tf.saved_model.load(
    export_dir, tags=None, options=None
)

與 SavedModel 關聯的簽名可作為函數使用：

imported = tf.saved_model.load(path)
f = imported.signatures["serving_default"]
print(f(x=tf.constant([[1.]])))

使用 tf.saved_model.save 導出的對象還具有分配給屬性的可跟蹤對象和函數：

exported = tf.train.Checkpoint(v=tf.Variable(3.))
exported.f = tf.function(
    lambda x:exported.v * x,
    input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32)])
tf.saved_model.save(exported, path)
imported = tf.saved_model.load(path)
assert 3. == imported.v.numpy()
assert 6. == imported.f(x=tf.constant(2.)).numpy()

加載 Keras 模型

Keras 模型是可跟蹤的，因此可以將它們保存到 SavedModel。 tf.saved_model.load 返回的對象不是 Keras 對象(即沒有 .fit , .predict 等方法)。一些屬性和函數仍然可用：.variables , .trainable_variables 和 .__call__。

model = tf.keras.Model(...)
tf.saved_model.save(model, path)
imported = tf.saved_model.load(path)
outputs = imported(inputs)

使用tf.keras.models.load_model 恢複 Keras 模型。

從 TensorFlow 1.x 導入 SavedModel

tf.estimator.Estimator 或 1.x SavedModel API 的 SavedModel 具有平麵圖，而不是 tf.function 對象。這些 SavedModel 將加載以下屬性：

• .signatures：將簽名名稱映射到函數的字典。

• .prune(feeds, fetches) ：一種允許您為新子圖提取函數的方法。這相當於在 TensorFlow 1.x 的會話中導入 SavedModel 並命名提要和提取。

  imported = tf.saved_model.load(path_to_v1_saved_model)
  pruned = imported.prune("x:0", "out:0")
  pruned(tf.ones([]))

  有關詳細信息，請參閱tf.compat.v1.wrap_function。

• .variables ：導入變量的列表。

• .graph ：整個導入的圖形。

• .restore(save_path) ：從 tf.compat.v1.Saver 保存的檢查點恢複變量的函數。

異步使用 SavedModels

當異步使用 SavedModels 時(生產者是一個單獨的進程)，SavedModel 目錄會在所有文件寫入之前出現，如果指向不完整的 SavedModel，tf.saved_model.load 將失敗。與其檢查目錄，不如檢查"saved_model_dir/saved_model.pb"。該文件作為最後一個tf.saved_model.save 文件操作以原子方式寫入。