致讀者:這是鐵人賽搶先版,還會有更多的陸續更新
在 Tensorflow 儲存深度學習模型有兩種模式,一是儲存模型所用的變數權重(tf.Variable
),被稱為 Checkpoint。而後者,則是包括模型的敘述等等可以讓深度學習模型在 Serving, mobile 或 embedded devices 或其他程式語言中使用,稱為 SavedModel。
前者需要模型的原始碼才能回覆模型,而後者則是與原始碼相互獨立。現在就將這兩種儲存模型的方式做介紹:
tf.keras.Model
API 儲存在 tf.keras
中,模型的儲存有三個模式。這些模式依據需寫入檔案的多寡和來分辨,如:全模型儲存(Whole-model saving),只有架構的儲存模式(Architecture-only saving)和只有參數的儲存模式(Weights-only saving)。現就全模型儲存需要存入哪一些檔案來做說明:
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 倒底存了什麼呢?
除了上面兩類的檔案外,還有一到兩個子資料夾 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 儲存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)