1. 程式人生 > >TensorFlow SavedModel儲存和載入模型

TensorFlow SavedModel儲存和載入模型

宣告:

SavedModel

如果你想儲存或恢復模型,我們推薦使用SaveModel. SaveModel是一種與語言無關,可恢復的密封式序列化格式。TensorFlow提供了多種與SavedModel互動的機制,如tf.saved_model API、Estimator API和CLI。

PS: 本文只介紹tf.saved_model API

儲存模型

使用SavedModel構建器tf.saved_model.builder.SavedModelBuilder. SavedModelBuilder類別提供了儲存多個MetaGraphDef的功能。MetaGraph是一種資料流圖,加上相關變數、資源和簽名。MetaGraphDef

MetaGraph的協議緩衝區的表示法。簽名是一組與圖有關的輸入和輸出。

此處的簽名並不是MetaGraphDef的tag,詳見SignatureDef

如果需要將資源儲存並寫入或複製到磁碟,則可以在首次新增 MetaGraphDef 時提供這些資源。如果多個 MetaGraphDef 與同名資源相關聯,則只保留首個版本。

必須使用使用者指定的標籤對每個新增到 SavedModel 的 MetaGraphDef 進行標註。這些標籤提供了一種方法來識別要載入和恢復的特定 MetaGraphDef,以及共享的變數和資源子集。這些標籤一般會標註 MetaGraphDef 的功能(例如服務或訓練),有時也會標註特定的硬體方面的資訊(如 GPU).

標準MetaGraphDef標籤

SavedModel為各種用例構建和載入TensorFlow圖提供了靈活性。對於最常見的用例,SavedModel的API在Python中提供了一組易於重複使用且便於在各種工具共享的變數。

tag Description
SEVING Tag for the serving graph
TRAINING Tag for the training graph
EVAL Tag for the eval graph. Not exported while the export logic is in contrib.
GPU Tag for the gpu graph
TPU Tag for the tpu graph

SavedModelBuilder

export_dir = ...
...
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
# Add the first MetaGraphDef
with tf.Session(graph=tf.Graph()) as sess:
  ...
  builder.add_meta_graph_and_variables(sess,
                                       [tag_constants.TRAINING],
                                       signature_def_map=foo_signatures,
                                       assets_collection=foo_assets)
...
# Add a second MetaGraphDef for inference.
with tf.Session(graph=tf.Graph()) as sess:
  ...
  builder.add_meta_graph([tag_constants.SERVING])
...
builder.save()

SaveModelBuilder: Builds the SavedModel protocol buffer and saved variables and assests.

The SavedModelBuilder class provides functionality to build a SavedModel protocol buffer. Specifically, this allows multiple meta graphs to be saved as part of a single language-neutral(語言無關的), while sharing variables and assets.

To build a SavedModel, the first meta graph must be saved with variables. Subsequent(後面的) meta graphs will simple be saved with their graph definitions. If assets need to be saved and written or copied to disk, they can be provided when the meta graph def is added. If multiple meta graph defs are associated an asset of the same name, only the first version is retained.

Each meta graph added to the SavedModel must be annotated with tags. The tags provide a means to identify the specific meta graph to load and restore, along with the shared set of variabels and assets.

Methods

__init__

__init__(export_dir)

add_meta_graph

add_meta_graph(tags, 
          signature_def_map=None, 
          asserts_collenction=None,
          legacy_init_op=None, 
          clear_devices=False,
          main_op=None, 
          strip_default_attrs=False)

Adds the current meta graph to the SavedModel.
Creates a Saver in current scope and uses the Saver to export the meta graph def. Invoking(呼叫) this API requires the add_meta_graph_and_variables() API to have been invoked before.

Args Description
tags The set of tags to annotate the meta graph def with
signature_def_map The map of signature defs to be added to the meta graph def.
assets_collection Assets collection to be saved with SavedModel. Note that this collection should be a subset of the assets saved as part of the first meta graph in the SavedModel.
legacy_init_op Legacy(遺留) support for op or group of ops to execute after the restore op upon a load.
clear_devices Set to true if the device info on the default graph should be cleared
main_op Op or group of ops to execute when the graph is loaded. Note that when the main_op is specified it is run after the restore op at load_time
strip_default_attrs Boolean. If True, default-valued attributes will be removed from the NodeDefs. 詳見Stripping Default-Valued Attributes.

add_meta_graph_and_variables

add_meta_graph_and_variables(
    sess, 
    tags, 
    signature_def_map=None,
    assets_collection=None, 
    legacy_init_op=None,
    clear_devices=False,
    main_op=None,
    strip_default_attrs=False)

Adds the current meta graph to the SaveModel and saved variables.
Creates a Saver to save the variables from the provided session. Exports the corresponding meta graph def. This function assumes that the variables to be saved have been initialized. For a given SavedModelBuilder, this API must be called exactly once and for the first meta graph to save. For subsequent meta graph defs to be adds, the add_meta_graph() API must be used.

save

save(as_text=False)

Writes a SavedModel protocol buffer to disk.
The fuction write the SavedModel protocol buffer to the export directory in serialized format.

引數as_text:Writes the SavedModel protocol buffer in text format to disk.
返回值:The Path to which the SavedModel protocol buffer was written.

恢復模型

使用SavedModel載入函式tf.saved_model.loader.load()為SavedModel提供載入和恢復功能。load指令需要以下資訊:

  • 要在其中恢復圖定義和變數的會話
  • 用於標識要載入的MetaGraphDef的標籤
  • SavedModel的位置(目錄).

載入後,作為特定MetaGraphDef的一部分,所提供的變數、資源和簽名子集將恢復到提供的會話中。

示例

export_dir = ...
...
with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, [tag_constants.TRAINING], export_dir)

load

tf.saved_model.loader.load(
    sess, 
    tags, 
    export_dir,
    **saver_kwargs)

Loads the model from a SavedModel as specified by tags.

Args Description
sess Session
tags Set of string tags to identify the required MetaGraphDef. These should correspond to the tags used when saving the variables using the SavedModel save API.
export_dir Directory in which the SaveModel protocol buffer and variables to be loaded and located.

返回值:The MetaGraphDef protocol buffer loaded in the provided session. This can be used to further extract signature-defs, collection-defs, etc.

SavedModel目錄的結構

當您以 SavedModel 格式儲存模型時,TensorFlow 會建立一個由以下子目錄和檔案組成的 SavedModel 目錄:

assets/
assets.extra/
variables/
    variables.data-?????-of-?????
    variables.index
saved_model.pb|saved_model.pbtxt

其中:

  • assets 是包含輔助(外部)檔案(如詞彙表)的子資料夾。資源被複制到 SavedModel 的位置,並且可以在載入特定的 MetaGraphDef 時讀取。
  • assets.extra 是一個子資料夾,其中較高級別的庫和使用者可以新增自己的資源,該資源與模型共存,但不會被圖載入。此子資料夾不由 SavedModel 庫管理。
  • variables 是包含 tf.train.Saver 的輸出的子資料夾。
  • saved_model.pbsaved_model.pbtxt 是 SavedModel 協議緩衝區。它包含作為 MetaGraphDef 協議緩衝區的圖定義。

單個 SavedModel 可以表示多個圖。在這種情況下,SavedModel 中所有圖共享一組檢查點(變數)和資源。例如,下圖顯示了一個包含三個 MetaGraphDef 的 SavedModel,它們三個都共享同一組檢查點和資源: