Categories
程式開發

TensorFlow 篇| TensorFlow 2.x 基於Keras 的模型保存及重建


TensorFlow 篇| TensorFlow 2.x 基於Keras 的模型保存及重建 1

「導語」模型訓練完成後一項十分重要的步驟是對模型信息進行持久化保存,以用於後續的再訓練以及線上Serving。保存後的模型文件除了個人使用外,還可以將其分享到TensorFlow Hub ,從而讓他人可以很方便地在該預訓練模型的基礎上進行再次開發與訓練。

Keras 模型概覽

一個常見的Keras 模型通常由以下幾個部分組成:

模型的結構或配置:它指定了模型中包含的層以及層與層之間的連接方式。

一系列權重矩陣的值:用於記錄模型的狀態。

優化器:用於優化損失函數,可以在模型編譯(compile) 時指定。

一系列損失以及指標:用於調控訓練過程,既包括在編譯時指定的損失和指標,也包括通過add_loss() 以及add_metric() 方法添加的損失和指標。

使用Keras API 可以將上述模型組成部分全部保存到磁盤,或者保存其中的部分。 Keras API 提供了以下三種選擇:

保存模型的全部內容:通常以TensorFlow SavedModel 格式或者Keras H5 格式進行保存,這也是最為常用的模型保存方式。

僅保存模型的結構:通常以json 文件的形式進行保存。

僅保存模型的權重矩陣的值:通常以numpy 數組的形式進行保存,一般會在模型再訓練或遷移學習時使用這種保存方式。

模型構建與訓練

在進行模型保存之前,我們需要先構建模型並進行訓練,使得模型損失和指標達到特定的狀態。考慮下面這個簡單的模型,它是使用Functional API 實現的一個3 層的深度神經網絡。

from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

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, name="predictions")(x)

model = keras.Model(inputs=inputs, outputs=outputs, name="3_layer_mlp")
model.summary()

x_train, y_train = (
np.random.random((60000, 784)),
np.random.randint(10, size=(60000, 1)),
)
x_test, y_test = (
np.random.random((10000, 784)),
np.random.randint(10, size=(10000, 1)),
)

model.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=keras.optimizers.RMSprop())
history = model.fit(x_train, y_train, batch_size=64, epochs=1)

# Save predictions for future checks
predictions = model.predict(x_test)

在使用隨機數據對該模型進行一輪的迭代訓練後,模型權重矩陣的值和優化器的狀態都發生了變化,此時我們可以對該模型進行保存操作了,當然,你也可以訓練之前進行模型保存,不過那樣做並沒有什麼實際意義。

保存整個模型

Keras 可以將模型的全部信息保存到一個文件目錄下,並且可以在之後使用該文件目錄重建模型,即使沒有模型的源碼也可以完成重建的過程。

保存後的模型文件中,包含了以下信息:

模型的結構。

模型的權重矩陣,這是模型在訓練過程中學到的信息。

模型的編譯(compile) 信息。

模型的優化器狀態信息,該狀態信息可以使你能夠在上次訓練終止的位置繼續進行訓練。

在Keras 中,可以使用model.save() 方法或者tf.keras.models.save_model() 方法對模型進行保存,使用tf.keras.models.load_model() 方法來對保存的模型文件進行加載並重建模型。

模型文件保存的格式可以有2 種,一種為TensorFlow SavedModel 格式,另一種是Keras H5 格式,官方推薦使用SavedModel 格式進行模型保存,它也是model.save() 方法默認使用的保存格式。可以通過設置save() 方法中的參數format=”h5″ 或者指定save() 方法中的文件名後綴為.h5 或.keras 來切換為使用H5 格式進行模型保存。

SavedModel 格式保存

使用SavedModel 格式進行模型保存及重建的代碼如下所示:

# Export the model to a SavedModel
model.save('path_to_saved_model')

# Recreate the exact same model
new_model = keras.models.load_model('path_to_saved_model')

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer state is preserved as well:
# you can resume training where you left off.
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

執行上述代碼會生成一個名為path_to_saved_model 的目錄,模型的全部信息會保存在該目錄下,目錄結構如下所示:

.
├── assets
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index

其中assets 是一個可選的目錄,用於存放模型運行所需的輔助文件,比如字典文件等。 variables 目錄下存放的是模型權重的檢查點文件,模型的所有權重信息均保存在該目錄下。 saved_model.pb 文件中包含了模型的結構以及訓練的配置信息如優化器,損失以及指標等信息。

H5 格式保存

Keras 也支持使用H5 格式進行模型保存及重建,保存後的文件中包含有模型的結構,權重矩陣的值以及編譯信息等,相比於SavedModel 格式, H5 格式是一個更為輕量的替代方案。

使用H5 格式進行模型保存及重建的代碼如下所示:

# Save the model
model.save('path_to_my_model.h5')

# Recreate the exact same model purely from the file
new_model = keras.models.load_model('path_to_my_model.h5')

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer state is preserved as well:
# you can resume training where you left off.
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

但使用H5 格式有它本身的局限性,相比於SavedModel , H5 格式的模型文件中不能保存以下兩種信息:

外部添加的損失以及指標信息:通過model.add_loss() 方法以及model.add_metric() 方法添加的損失和指標信息不會保存到H5 文件中。如果你想加載模型後繼續在訓練中使用它們,需要重新通過上述方法進行添加。需要注意的是在Keras 層內通過self.add_loss() 或self.add_metric() 添加的損失和指標是會被H5 文件自動保存的。

自定義對象的計算圖信息:比如自定義的層(layers) 不會被保存到H5 文件中。此時如果使用H5 文件直接進行模型加載,會報ValueError: Unknown layer 錯誤。

因此在保存整個模型時,尤其是對於自定義(Subclassed) 的模型,我們最好使用SavedModel 格式來進行模型保存以及重建。

另外從代碼裡可以看到,由於model.save() 方法保存了模型的全部信息,所以在重建模型後,可以在原有模型的基礎上繼續進行訓練,而無需進行其它額外的設置。

僅保存模型結構

有時我們可能只對模型的結構感興趣,而不想保存模型的權重值以及優化器的狀態等信息。在這種情況下,我們可以藉助於模型的配置方法(config) 來對模型的結構進行保存和重建。

順序/功能

對於Sequential 模型和Functional API 模型,因為它們大多是由預定義的層(layers) 組成的,所以它們的配置信息都是以結構化​​的形式存在的,可以很方便地執行模型結構的保存操作。我們可以使用模型的get_config() 方法來獲取模型的配置,然後通過from_config(config) 方法來重建模型。

對於最開始介紹的Functional API 的模型而言,其結構保存與重建的代碼如下所示:

config = model.get_config()
new_model = keras.Model.from_config(config)

# Note that the model state is not preserved! We only saved the architecture.
new_predictions = new_model.predict(x_test)
assert abs(np.sum(predictions - new_predictions)) > 0.

其中get_config() 方法返回的結果是一個python 字典,它包含了該模型結構的全部信息,因此可以很方便地進行模型重建。 python 字典的內容如下所示:

{
'name':
'3_layer_mlp',
'layers': [{
'class_name': 'InputLayer',
'config': {
'batch_input_shape': (None, 784),
'dtype': 'float32',
'sparse': False,
'ragged': False,
'name': 'digits'
},
'name': 'digits',
'inbound_nodes': []
}, {
'class_name': 'Dense',
'config': {
'name': 'dense_1',
'trainable': True,
'dtype': 'float32',
'units': 64,
'activation': 'relu',
'use_bias': True,
'kernel_initializer': {
'class_name': 'GlorotUniform',
'config': {
'seed': None
}
},
'bias_initializer': {
'class_name': 'Zeros',
'config': {}
},
'kernel_regularizer': None,
'bias_regularizer': None,
'activity_regularizer': None,
'kernel_constraint': None,
'bias_constraint': None
},
'name': 'dense_1',
'inbound_nodes': [[['digits', 0, 0, {}]]]
}, {
'class_name': 'Dense',
'config': {
'name': 'dense_2',
'trainable': True,
'dtype': 'float32',
'units': 64,
'activation': 'relu',
'use_bias': True,
'kernel_initializer': {
'class_name': 'GlorotUniform',
'config': {
'seed': None
}
},
'bias_initializer': {
'class_name': 'Zeros',
'config': {}
},
'kernel_regularizer': None,
'bias_regularizer': None,
'activity_regularizer': None,
'kernel_constraint': None,
'bias_constraint': None
},
'name': 'dense_2',
'inbound_nodes': [[['dense_1', 0, 0, {}]]]
}, {
'class_name': 'Dense',
'config': {
'name': 'predictions',
'trainable': True,
'dtype': 'float32',
'units': 10,
'activation': 'linear',
'use_bias': True,
'kernel_initializer': {
'class_name': 'GlorotUniform',
'config': {
'seed': None
}
},
'bias_initializer': {
'class_name': 'Zeros',
'config': {}
},
'kernel_regularizer': None,
'bias_regularizer': None,
'activity_regularizer': None,
'kernel_constraint': None,
'bias_constraint': None
},
'name': 'predictions',
'inbound_nodes': [[['dense_2', 0, 0, {}]]]
}],
'input_layers': [['digits', 0, 0]],
'output_layers': [['predictions', 0, 0]]
}

對於Sequential 模型而言,模型結構保存與重建的示例代碼如下所示:

from tensorflow import keras

model = keras.Sequential([keras.Input((32, )), keras.layers.Dense(1)])
config = model.get_config()
new_model = keras.Sequential.from_config(config)
new_model.summary()

為了方便模型結構的持久化存儲,可以直接使用模型的to_json() 方法將模型的信息序列化為json 格式然後保存到本地,重建模型時可以使用from_json() 方法解析該json 文件內容以完成重建工作。示例代碼如下所示:

json_config = model.to_json()
new_model = keras.models.model_from_json(json_config)

需要注意的是因為只保存了模型的結構,所以在模型重建後,新模型並不包含之前模型的編譯信息和狀態信息,因此需要重新對新模型進行編譯以進行後續的訓練操作。

Subclassed 模型和層

Subclassed 模型的結構信息是在__init__ 和call 方法中定義的,它們會被視為python 字節碼而無法將其序列化為json 格式。截至本文完成前,還不存在一種方法可以直接對自定義模型(繼承自keras.Model)的結構信息進行保存和重建操作。

而對於Subclassed 層(繼承自keras.Layer )來說是可以進行結構保存和重建的,我們需要重寫其get_config() 方法和from_config() (可選的)方法來達成這一目標。

其中get_config() 方法需要返回一個可被json 序列化的字典,以便適配Keras 的結構以及模型保存的接口。 from_config(config) 類方法需要根據config 參數返回一個新的層對象。

定義好層的config 方法之後,我們可以使用serialize 和deserialize 方法來序列化(保存)和反序列化(重建)層的結構信息,注意,在重建層時需要以custom_objects 方式提供自定義層的信息。示例代碼如下所示:

from tensorflow import keras

class ThreeLayerMLP(keras.layers.Layer):
def __init__(self, hidden_units):
super().__init__()
self.hidden_units = hidden_units
self.dense_layers = [keras.layers.Dense(u) for u in hidden_units]

def call(self, inputs):
x = inputs
for layer in self.dense_layers:
x = layer(x)
return x

def get_config(self):
return {"hidden_units": self.hidden_units}

layer = ThreeLayerMLP([64, 64, 10])
serialized_layer = keras.layers.serialize(layer)
print(serialized_layer)
new_layer = keras.layers.deserialize(
serialized_layer,
custom_objects={"ThreeLayerMLP": ThreeLayerMLP},
)

其中serialize 方法返回的結果如下所示:

{'class_name': 'ThreeLayerMLP', 'config': {'hidden_units': ListWrapper([64, 64, 10])}}

對於使用自定義層和自定義函數構造的Functional API 模型,既可以使用上述的serialize/deserialize 方法來保存和重建模型結構,也可以通過get_config/keras.Model.from_config 方法來完成該保存和重建操作,不過需要藉助於keras.utils.custom_object_scope 方法來指明自定義的對象。示例代碼如下所示:

def custom_activation(x):
return tf.nn.tanh(x)**2

inputs = keras.Input((64, ))
x = layer(inputs)
outputs = keras.layers.Activation(custom_activation)(x)
model = keras.Model(inputs, outputs)

config = model.get_config()
print(config)
custom_objects = {
"ThreeLayerMLP": ThreeLayerMLP,
"custom_activation": custom_activation
}
with keras.utils.custom_object_scope(custom_objects):
new_model = keras.Model.from_config(config)
new_model.summary()

另外對於上述的Functional API 模型,還可以直接在內存中進行模型的拷貝,這與獲取配置然後再通過配置進行模型重建的流程基本一致。示例代碼如下所示:

with keras.utils.custom_object_scope(custom_objects):
new_model = keras.models.clone_model(model)
new_model.summary()

僅保存模型權重

除了可以選擇只保存模型的結構信息外,你也可以只保存模型的權重信息。它在以下場景中會比較有用:

你只需要使用模型進行預測:在這種情況下,你不需要繼續訓練現有模型,此時就沒有必要保存模型的編譯信息以及優化器的狀態信息。

你正在進行遷移學習:在這種情況下,你會使用現有模型的狀態(權重矩陣)來訓練一個新的模型,因此也不需要現有模型的編譯信息。

在Keras 中,可以通過模型的get_weights() 和set_weights() 方法來獲取和設置權重矩陣的值。代碼如下所示:

weights = model.get_weights() # Retrieves the state of the model.
model.set_weights(weights) # Sets the state of the model.

模型的get_weights() 方法返回的是一個numpy array 組成的list ,裡面記錄了模型的所有權重矩陣的信息,後面可以將保存的權重信息作為set_weights(weights) 方法的參數來設置新模型的狀態,也可以將這項操作理解為模型間的權重遷移。

除了可以在同類模型間進行權重遷移外,具有兼容體系結構的模型之間同樣也可以進行權重遷移操作,比如我們可以先獲取前面Functional 模型的權重矩陣信息,然後將該權重矩陣的值通過set_weights( ) 方法遷移到相應的Subclassed 模型中。不過需要注意的是,兩個模型中的權重矩陣的順序,數量以及大小要保持一致。示例代碼如下所示:

from tensorflow import keras
from tensorflow.keras import layers

class ThreeLayerMLP(keras.Model):
def __init__(self, hidden_units):
super().__init__()
self.hidden_units = hidden_units
self.dense_layers = [layers.Dense(u) for u in hidden_units]

def call(self, inputs):
x = inputs
for layer in self.dense_layers:
x = layer(x)
return x

def get_config(self):
return {"hidden_units": self.hidden_units}

subclassed_model = ThreeLayerMLP([64, 64, 10])
subclassed_model(tf.ones((1, 784)))
subclassed_model.set_weights(functional_model.get_weights())

assert len(functional_model.weights) == len(subclassed_model.weights)
for a, b in zip(functional_model.weights, subclassed_model.weights):
np.testing.assert_allclose(a.numpy(), b.numpy())

另外,對於無狀態的層,由於其不包含權重矩陣,所以它不會改變模型中權重矩陣的順序以及數量,因此即使不同的模型間存在額外的無狀態層(如Dropout 層),它們也可以擁有兼容的模型結構,同樣可以使用上述方式進行權重矩陣的遷移操作。

如果要將權重矩陣信息持久化到本地,可以使用save_weights(fpath) 方法進行本地保存,後面可以使用load_weights(fpath) 方法從本地文件中加載權重信息來恢復模型的狀態。示例代碼如下:

# Save weights to disk
model.save_weights('path_to_my_weights.h5')
new_model = ThreeLayerMLP([64, 64, 10])
new_model(tf.ones((1, 784)))
new_model.load_weights('path_to_my_weights.h5')

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer was not preserved.
new_model.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=keras.optimizers.RMSprop())
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

注意,save_weights 方法既可以保存為Keras HDF5 格式的文件,也可以保存為TensorFlow Checkpoint 格式的文件。與model.save 方法類似,它也包含一個參數save_format ,同樣有兩種取值。默認為tf 表示使用Checkpoint 格式保存, h5 表示使用HDF5 格式保存。如果不明確指定,它會根據文件的後綴名去推測,以.h5 或.hdf5 結尾的文件名會以HDF5 格式進行保存。

可以結合使用get_config()/from_config() 和get_weights()/set_weights() 來重建模型並恢復原有的狀態。但是與使用model.save() 不同,這種方式因為不能保存模型的訓練配置以及優化器的狀態,所以在繼續訓練時,需要重新調用模型的compile 方法來指定訓練過程所需的一些配置。代碼如下所示:

config = model.get_config()
weights = model.get_weights()

new_model = keras.Model.from_config(config)
new_model.set_weights(weights)

# Check that the state is preserved
new_predictions = new_model.predict(x_test)
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

# Note that the optimizer was not preserved,
# so the model should be compiled anew before training
# (and the optimizer will start from a blank state).
new_model.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=keras.optimizers.RMSprop())
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

TensorFlow集線器

TensorFlow Hub 是一個用於可重用機器學習的模型存儲庫和函數庫。 TensorFlow Hub 中的模型可以在不同的學習任務中被重用(即遷移學習),它能讓我們在訓練時僅使用很小的數據集就能取得不錯的效果,同時還能提升模型的泛化能力以及加快模型的學習速度。

當我們在實現一個較為常見的模型時,可以優先考慮使用TensorFlow Hub 中已有的預訓練模型,並在該模型基礎上繼續建模與訓練,以減少程序的代碼量和模型的訓練時間。

TensorFlow 2.x 推薦使用SavedModel 格式的文件在TensorFlow Hub 上分享預訓練的模型。 tensorflow_hub 函數庫提供了一個類hub.KerasLayer ,它可以從遠程或本地的SavedModel 文件中讀取模型信息並以Keras Layer 的方式重建模型,重建後的模型中包含有預訓練的權重矩陣信息。

使用hub.KerasLayer 加載模型並預測的代碼如下所示:

import tensorflow_hub as hub

new_layer = hub.KerasLayer(
'path_to_my_model',
trainable=True,
input_shape=[784],
dtype=tf.float32,
)

# Check that the state is preserved
new_predictions = new_layer(x_test.astype('float32'))
np.testing.assert_allclose(predictions, new_predictions, rtol=1e-6, atol=1e-6)

在預測時,要注意保存的模型文件中權重矩陣的數據類型與輸入數據的類型相匹配。在TensorFlow 2.x 中,模型文件的數據類型由tf.keras.backend.floatx() 獲得,默認為float32。可以通過tf.keras.backend.set_floatx(‘float64’) 對模型的數據類型進行修改或者將輸入數據進行類型轉化從而避免因類型不匹配而出錯。

使用該方法加載模型與使用keras.models.load_model 加載模型相比有幾點不同:

其一,它本質上是一個Keras Layer,因此不具有模型的一些方法,比如summary 方法。

其二,其內部結構已經被隱藏了,因此不能查看其具體的組成結構(layers) 。

其三,在訓練時,需要將其作為模型的一層,並創建一個新的模型才能繼續訓練。

在Keras 代碼中使用hub.KerasLayer 進行訓練的代碼如下:

new_layer = hub.KerasLayer(
'path_to_my_model',
trainable=True,
input_shape=[784],
dtype=tf.float32,
)

new_model = tf.keras.Sequential([new_layer])
new_model.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=keras.optimizers.RMSprop())
new_model.fit(x_train, y_train, batch_size=64, epochs=1)

從代碼中可以看到,我們可以像使用Keras Layer 一樣使用hub.KerasLayer 。默認情況下, hub.KerasLayer 中的權重矩陣的值是不可訓練的,即在新的訓練過程中保持不變,如果需要對原來的模型做微調,可以將trainable 參數設置為True 。

另外,如果需要將微調後的模型重新導出,則可以使用上面介紹的各種方法如model.save() 或tf.keras.models.save_model() 等按需導出。

如果你想將自己的預訓練模型分享到TensorFlow Hub ,則首先需要將模型導出為SavedModel 格式的文件,然後再推送到線上進行模型分享。在導出模型文件時需要注意,如果是要分享整個模型,在調用模型的save 方法時需要將include_optimizer 參數設置為False ,也就是說分享的模型中不能包含優化器的狀態信息;如果只是要分享模型的某一部分,則可以在構建模型前,將待分享的部分抽離出來,並在訓練完成後保存即可。示例代碼如下所示:

piece_to_share = tf.keras.Model(...)
full_model = tf.keras.Sequential([piece_to_share, ...])
full_model.fit(...)
piece_to_share.save(..., include_optimizer=False)

參考資料

保存和加載Keras模型TensorFlow 2中的TF Hub中的SavedModels