TF2.0模型保存

概述

这是TF2.0入门笔记【TF2.0模型创建、TF2.0模型训练、TF2.0模型保存】中最后一篇【TF2.0模型保存】，前两篇分别介绍了模型的创建与训练，本篇将介绍模型的保存以及加载。

模型保存及加载

事实上，tensorflow1.x和tensorflow2.0的模型保存略有差别，tensorflow1.x还有一个freeze_graph的步骤…。但是这里就不做讨论了，本篇仅以TF2.0作为介绍。

import tensorflow as tf
import os
from tensorflow.keras.layers import Dense, Flatten, Conv2D
from tensorflow.keras import Sequential
from tensorflow.keras.models import load_model

model = Sequential()

model.add(Conv2D(32, 3, activation='relu', padding='same', input_shape=(28, 28, 1)))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dense(10, activation='softmax'))

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

x_train = x_train[..., tf.newaxis]
x_test  = x_test[..., tf.newaxis]

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

保存为.h5格式

1、保存权重

保存方式可以通过callback保存，也可以在最后通过调用save_weights方法保存
1.1、通过callback保存

model_path = "./model"
model_name = "weight_callback.h5"
if not os.path.exists(model_path):
    os.mkdir(model_path)

model_file = os.path.join(model_path, model_name)
callback = tf.keras.callbacks.ModelCheckpoint(filepath=model_file,
                                              save_weights_only=True,
                                              verbose=1)

model.fit(x_train, y_train, epochs=1, callbacks=[callback])

1.2、在最后调用save_weights方法保存

model.save_weights('./model/weight_save.h5')

运行上面的两个代码，我们会得到以下两个权重文件：

2、加载权重

接下来我们加载一下看看它们测试结果是不是一样的。
注意：由于是我们保存的是模型的参数，所以我们加载的时候需要需要重新定义模型结构

model = Sequential()

model.add(Conv2D(32, 3, activation='relu', padding='same', input_shape=(28, 28, 1)))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dense(10, activation='softmax'))

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

2.1、加载使用callback方法保存的权重

model.load_weights('./model/weight_callback.h5')
model.evaluate(x_test,  y_test, verbose=2)

运行输出：

10000/10000 - 2s - loss: 0.0644 - accuracy: 0.9781
[0.06437707488704472, 0.9781]

2.2、加载使用save_weights方法保存的权重

model.load_weights('./model/weight_save.h5')
model.evaluate(x_test,  y_test, verbose=2)

运行输出：

10000/10000 - 2s - loss: 0.0644 - accuracy: 0.9781
[0.06437707488704472, 0.9781]

分析：我们可以看到，结果是一样的

3、保存整个模型

保存整个模型跟保存权重一样，也是两种方法。
1、通过callback保存
这里把save_weights_only改成False就可以了

model_path = "./model"
model_name = "model_callback.h5"
if not os.path.exists(model_path):
    os.mkdir(model_path)

model_file = os.path.join(model_path, model_name)
callback = tf.keras.callbacks.ModelCheckpoint(filepath=model_file,
                                              save_weights_only=False,
                                              verbose=1)

model.fit(x_train, y_train, epochs=1, callbacks=[callback])

2、在最后调用save方法保存

model.save('./model/model_save.h5')

运行上面的两个代码，我们会得到以下两个模型文件：

4、加载模型

由于这次我们保存的是整个模型（包括参数），所以加载的时候不需要重新定义模型结构，只需要通过load_model读取文件就可以了
4.1、加载使用callback方法保存的模型

model = load_model('./model/model_callback.h5')
model.evaluate(x_test,  y_test, verbose=2)

运行输出：

10000/10000 - 2s - loss: 0.0602 - accuracy: 0.9814
[0.0602496856124606, 0.9814]

4.2、加载使用save方法保存的模型

model = load_model('./model/model_save.h5')
model.evaluate(x_test,  y_test, verbose=2)

运行输出：

10000/10000 - 2s - loss: 0.0602 - accuracy: 0.9814
[0.0602496856124606, 0.9814]

保存为tf格式

事实上，TF2.0已经不推荐或者说不支持将参数保存为.ckpt格式了，这里演示的也不能说是.ckpt格式，官方文档称之为’tf’格式，如果你不指定是’h5’格式，那它就是’tf’格式。（个人不推荐使用这种方法，事实上这里我也不太懂）

model_path = "./model"
model_name = "weight_callback.ckpt"
if not os.path.exists(model_path):
    os.mkdir(model_path)

model_file = os.path.join(model_path, model_name)
callback = tf.keras.callbacks.ModelCheckpoint(filepath=model_file,
                                              save_weights_only=True,
                                              verbose=1)

model.fit(x_train, y_train, epochs=1, callbacks=[callback])

运行得到以下文件：
这里只要你文件名的后缀不是.h5（如果后缀不是.h5，那么就会保存成’tf’格式，这时后缀已经变得不重要了，你甚至可以不要后缀），那么保存之后你就会得到以下三个文件。

model.load_weights('./model/weight_callback.ckpt')
model.evaluate(x_test,  y_test, verbose=2)

运行输出：

10000/10000 - 2s - loss: 0.0637 - accuracy: 0.9789
[0.06373853980554267, 0.9789]

保存为.pb格式（Saved Model）

1、保存成Saved Model

这里也仍然有两种方式将模型保存为Saved Model格式。
1.1、通过callback保存
想要通过callback保存成Saved Model格式，必须要传递的是一个文件夹，而不是文件名，而且save_weights_only=False这个参数一定要设成False，不然就会保存为之前说的’tf’格式，注意：这两个条件缺一不可。

model_path = "./model"

if not os.path.exists(model_path):
    os.mkdir(model_path)

callback = tf.keras.callbacks.ModelCheckpoint(filepath=model_path,
                                              save_weights_only=False,
                                              verbose=1)

1.2、在最后用tf.saved_model.save方法保存

tf.saved_model.save(model, './model')

运行上面的两个代码，我们会得到以下文件（两种方法运行结果一样）：

2、加载Saved Model

这里你仍然可以简单地通过load_model去加载Saved Model，Saved Model也是保存的是整个模型，但是Saved Model还有更多的用处，不过这里就不深入介绍了（以后要是出TF2.0模型部署会详细介绍）。

model = load_model('./model/')
model.evaluate(x_test,  y_test, verbose=2)

运行输出：

10000/10000 - 2s - loss: 0.0663 - accuracy: 0.9790
[0.06630929337199777, 0.979]

完结撒花

TF2.0入门笔记【TF2.0模型创建、TF2.0模型训练、TF2.0模型保存】三部曲