Day 21: Tensorflow 2.0: 再造訪 Model Persistence

11th鐵人賽

renewang

2019-10-06 23:37:38

12740 瀏覽

致讀者：這是鐵人賽搶先版，還會有更多的陸續更新

在 Tensorflow 儲存深度學習模型有兩種模式，一是儲存模型所用的變數權重（tf.Variable），被稱為 Checkpoint。而後者，則是包括模型的敘述等等可以讓深度學習模型在 Serving, mobile 或 embedded devices 或其他程式語言中使用，稱為 SavedModel。

前者需要模型的原始碼才能回覆模型，而後者則是與原始碼相互獨立。現在就將這兩種儲存模型的方式做介紹：

SavedModel

使用 `tf.keras.Model` API 儲存

在 tf.keras 中，模型的儲存有三個模式。這些模式依據需寫入檔案的多寡和來分辨，如：全模型儲存（Whole-model saving），只有架構的儲存模式（Architecture-only saving）和只有參數的儲存模式（Weights-only saving）。現就全模型儲存需要存入哪一些檔案來做說明：

模型的架構（The model's architecture）
模型的參數值（The model's weight value）
模型的訓練組態（The model's training config）
最佳化演算法和它的狀態（The optimizer and its state）
要執行全模型儲存，在 tf.keras.Model 物件中，可以透過呼叫 model.save() 來完成寫出以及 model.load()來完成讀入。至於如何使用，可以根據下面的程式碼，將模型寫出為 SavedModel 格式。

# 輸出 SavedModel 的格式來輸出模型
model.save('path_to_saved_model', save_format='tf')

# 載入模型並建立模型物件（不需要模型的類別宣告原始碼，即可完成）
new_model = keras.models.load_model('path_to_saved_model')

至於 SavedModel 倒底存了什麼呢？

包括參數值的 Checkpoint 檔案，這些檔案通常位於 variable 子資料夾。關於 Checkpoint，我們會流到下段文章再介紹。
包括計算圖的 protobuf 檔案，通常命名為 saved_model.pb

除了上面兩類的檔案外，還有一到兩個子資料夾 assets 或 assets.extra。這兩個資料夾都是儲存一些輔助重建計算圖的檔案，如字匯集和第三方軟體客製化 SavedModel 的檔案。

除了全模型儲存的方式，另一種方式則是只有模型架構的儲存方式。這個方式只需要能回復模型架構的資訊，而不會使用參數值，或最佳化演算法的狀態。要取得架構資訊，可以呼叫 get_config 或 to_json，前者是回傳 dict 物件，而後者是將模型架構資訊以 json 的格式寫出。而若要使用以載入的 cofing 或 json format 的模型資料，則呼叫 tf.keras.Model.from_config 或 from_json 的方式來實例化模型。更多關於使用的細節可以見下面程式碼：

# dict in-memory 版本
config = model.get_config()
# 不需要模型的程式碼，直接用類別方法載入模型細節
reinitialized_model = tf.keras.Model.from_config(config)

# json 版本
json_config = model.to_json()
reinitialized_model = keras.models.model_from_json(json_config)

在之前的範例中，我們都假設 tf.keras.Model 是由 functional API 所建構的，而不是透過繼承的方式來建立一個子類別後實例化。事實上，若是仰賴繼承的方式建立的模型，在序列化模型時都需要考慮原始碼是否能夠寫入檔案，以及是否寫成 python byte code 的方式（使用 python pickle）。

使用 `tf.estimator` API 儲存

Checkpoint

使用 `tf.keras` training API

寫入 Checkpoint 格式，最方便的方法即是使用 tf.keras training API，並呼叫 tf.keras.Model.save_weights。這個tf.keras.Model.save_weights函式只儲存權重（weights-only）。而 tf.keras.Model 可以藉由 functional API 建立，若是使用繼承的方式來建立模型的則有些限制，我們將會在最後提到。
tf.keras.Model 有 get_weights 和 set_weights。呼叫的方法大概如這段程式碼：

from tensorflow import keras
from tensorflow.keras import layers

# 使用 Functional API 來建立 keras.Model
inputs = keras.Input(shape=(784,), name='digits')
x = layers.Dense(64, activation='relu', name='dense_1')(inputs)
x = layers.Dense(64, activation='relu', name='dense_2')(x)
outputs = layers.Dense(10, activation='softmax', name='predictions')(x)
model = keras.Model(inputs=inputs, outputs=outputs, name='3_layer_mlp')

# 訓練 ....

model.save_weights('path_to_my_tf_checkpoint', save_format='tf')
# model 的原始碼必須存在
model.load_weights('path_to_my_tf_checkpoint')
weights = model.get_weights()  # Retrieves the state of the model.
model.set_weights(weights)  # Sets the state of the model.

save_format 必須要填入 'tf' 才會是 Checkoint 格式，另一種格式則是 hdf5。

使用 `tf.estimator` 物件導向的儲存方式

若建立 tf.estimator 物件，想要將 tf.estimator使用的參數權重儲存，只要在 tf.estimator的 constructor 傳入欲儲存的

def model_fn(features, labels, mode):
  net = Net()
  opt = tf.keras.optimizers.Adam(0.1)
  ckpt = tf.train.Checkpoint(step=tf_compat.train.get_global_step(),
                             optimizer=opt, net=net)
  with tf.GradientTape() as tape:
    output = net(features['x'])
    loss = tf.reduce_mean(tf.abs(output - features['y']))
  variables = net.trainable_variables
  gradients = tape.gradient(loss, variables)
  return tf.estimator.EstimatorSpec(
    mode,
    loss=loss,
    train_op=tf.group(opt.apply_gradients(zip(gradients, variables)),
                      ckpt.step.assign_add(1)),
    # Checkpoint 將會以物件為基礎的格式
    scaffold=tf_compat.train.Scaffold(saver=ckpt))

tf.keras.backend.clear_session()
est = tf.estimator.Estimator(model_fn, './tf_estimator_example/')
est.train(toy_dataset, steps=10)
# 訓練...
opt = tf.keras.optimizers.Adam(0.1) # 建立一個一樣的 optimizer
net = Net()  ＃需要模型的原始碼
ckpt = tf.train.Checkpoint( # 回覆
  step=tf.Variable(1, dtype=tf.int64), optimizer=opt, net=net)
ckpt.restore(tf.train.latest_checkpoint('./tf_estimator_example/'))
ckpt.step.numpy()  # From est.train(..., steps=10)

手動儲存

要儲存為 Checkpoint 首先要生成 tf.train.Checkpoint 物件，如下面的程式碼：

# 生成 tf.train.Checkpoint 物件，傳入 optimizer，模型 net 和 global step 變數（初始值為 1）
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net)
# 生成一個 tf.train.CheckpointManager，傳入 tf.train.Checkpoint 物件，路徑和參數
manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)

上面的程式碼，模型 net 是繼承 tf.keras.Model 的新類別（程式碼如下）

class Net(tf.keras.Model):
  """A simple linear model."""

  def __init__(self):
    super(Net, self).__init__()
    self.l1 = tf.keras.layers.Dense(5)

  def call(self, x):
    return self.l1(x)