模型的保存和加載
訓練一個相對復雜的模型很有可能需要一段時間,如果是在專門的服務器或計算資源上進行訓練那放那里跑就行了。但是如果是在自己的小電腦上跑,就干等着,就可能這段時間電腦都用不了。萬一期間要做個其他實驗,或者單純打個游戲放松下就難受了。
好在TensorFlow提供了訓練期間和訓練后對模型保存的方法。也就是說,你可以隨時暫停一下,然后隨時恢復繼續訓練,甚至別人訓練了一半,你可以拿它們的權重繼續訓練,避免耗費長時間在訓練上。
保存TensorFlow模型有許多不同的方法,取決於你使用的API。因為tf2.0內置了keras這里基本使用keras的方式來完成模型的保存和加載。
- windows10
- python==3.6
- tensorflow==2.0
准備模型
創建一個簡單的MNIST手寫數字識別模型
獲取數據集
為了訓練速度,可以只使用前1000個樣本來進行訓練和測試
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_labels = train_labels[:1000]
test_labels = test_labels[:1000]
train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0
定義模型
這里就直接FC結構不用卷積了
# Define a simple sequential model
def create_model():
model = tf.keras.models.Sequential([
keras.layers.Dense(512, activation='relu', input_shape=(784,)),
keras.layers.Dropout(0.2),
keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
return model
# Create a basic model instance
model = create_model()
# Display the model's architecture
model.summary()
可以看到如下輸出
Model: "sequential"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
dense (Dense) (None, 512) 401920
_________________________________________________________________
dropout (Dropout) (None, 512) 0
_________________________________________________________________
dense_1 (Dense) (None, 10) 5130
=================================================================
Total params: 407,050
Trainable params: 407,050
Non-trainable params: 0
_________________________________________________________________
使用回調保存模型
tf.keras.callbacks.ModelCheckpoint
回調可以讓我們在訓練期間和訓練結束時保存模型
創建回調實例
checkpoint_path = "training_1/cp.ckpt" # your path to save model
# checkpoint_dir = os.path.dirname(checkpoint_path) # training_1
# Create a callback that saves the model's weights
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
save_weights_only=True,
verbose=1)
# Train the model with the new callback
model.fit(train_images,
train_labels,
epochs=10,
validation_data=(test_images,test_labels),
callbacks=[cp_callback]) # Pass callback to training
這個回調會在每個epoch后保存(會覆蓋舊的)checkpoint。
使用保存的模型
因為我們只保存了模型的權值,所以在恢復的時候需要先創建一個同樣結構的模型,然后載入我們保存的權值。
# Create a basic model instance
model = create_model()
# Evaluate the model
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Untrained model, accuracy: {:5.2f}%".format(100*acc))
我們還沒有對模型進行訓練,所以現在權值是全隨機的,Accuracy也不會很好。
1000/1 - 0s - loss: 2.3598 - accuracy: 0.1360
Untrained model, accuracy: 13.60%
接下來從保存的checkpoint中加載權值再測試
# Loads the weights
model.load_weights(checkpoint_path)
# Re-evaluate the model
loss,acc = model.evaluate(test_images, test_labels, verbose=2)
print("Restored model, accuracy: {:5.2f}%".format(100*acc))
之前保存的權值被加載到新創建的模型中,Accuracy變成了之前訓練保存的。
1000/1 - 0s - loss: 0.3937 - accuracy: 0.8650
Restored model, accuracy: 86.50%
當然tf.keras.callbacks.ModelCheckpoint有很多參數可以設置,比如保存的間隔,或者你可以讓他不覆蓋,依次保存為不同名字的模型,這樣可以在多個checkpoint中選擇沒有過擬合的
手動保存模型
手動保存模型的方式很簡單,使用tf.keras.Model.save_weights
方法即可。當然前提是你的模型是一個tf.keras.Model
類型。
# Save the weights
model.save_weights('./checkpoints/my_checkpoint')
# Create a new model instance
model = create_model()
# Restore the weights
model.load_weights('./checkpoints/my_checkpoint')
保存整個模型
之前我們保存的一直都是模型的權值,而沒有保存模型的結構,所以我們恢復的時候需要創建一個相同結構的模型,然后使用tf.keras.Model.load_weights
來加載權重。
保存整個模型可以讓我們使用時不需要手動寫模型結構,直接從保存的模型里提取出模型。
安裝依賴
如果要使用HDF5格式來保存整個模型需要安裝如下依賴
pip install -q pyyaml h5py
HDF5格式
保存為HDF5格式只需要使用tf.keras.Model.save()
,且文件名后綴為.h5
即可
# Create and train a new model instance.
model = create_model()
model.fit(train_images, train_labels, epochs=5)
# Save the entire model to a HDF5 file.
# The '.h5' extension indicates that the model shuold be saved to HDF5.
model.save('my_model.h5')
# Or give a save_format arg
# model.save(filepath="", save_format="h5")
SavedModel格式
tf.keras.Model.save()
方法也可以用於保存SavedModel格式(這個格式是tf的格式,h5則是一個通用的格式)
# Create and train a new model instance.
model = create_model()
model.fit(train_images, train_labels, epochs=5)
# Save the entire model as a SavedModel.
!mkdir -p saved_model
model.save('saved_model/my_model')
# Or give a save_format arg
#model.save('saved_model/my_model', save_format='tf')
通過tf.keras.Model.save()
的參數save_format
可以調整保存的格式,tf2.0以下可能默認保存的是h5格式,tf2.0改為了tf格式。
加載模型
加載的時候就不需要先定義一個相同結構的模型,直接從保存的模型中提取出模型結構和權值
# Recreate the exact same model, including its weights and the optimizer
new_modelh5 = tf.keras.models.load_model('my_model.h5')
new_modeltf = tf.keras.models.load_model('saved_model/my_model')
# Show the model architecture
new_modelh5.summary()
new_modeltf.summary()
tf中tf.keras.models.load_model
可以恢復兩種格式保存的模型。
保存自定義的模型對象
tf中SavedModel
和HDF5
一大區別是,SavedModel
可以自動保存自定義的模型對象(有自定義層的模型),而HDF5
需要模型配置來保存自定義模型。
參考
[1] https://tensorflow.google.cn/tutorials/keras/save_and_load