详解tensorflow2.0的模型保存方法（一）

前言：tensorflow2.0中有很多种模型保存方法，为了做一个统一的比较，本文做一个归纳总结，加之自己在这上面踩过很多的坑，现在感觉做一个总结实在是有必要。因为tensorflow创建模型的方式很多，包括Sequence、Model、SubClass Model，我们可以保存所有的模型，也可以仅仅保存权重。

本文为系列文章第一篇，介绍各种模型保存方法，后面还会介绍checkpoint以及SavedModel的细节。

一、Sequence和Model模型的保存

1.1 保存整个模型为一个 h5 文件

由于是保存的整个模型，这个模型包括如下内容：

• The model's architecture，模型的结构
• The model's weight values (模型的权值)
• The model's training config (模型的配置，即我们通过compile编译模型的一些信息，如优化器，损失函数等), 
• The optimizer and its state, (优化器的状态信息，我们可以接着之前的训练继续训练)

# 保存模型
model.save('path_to_my_model.h5')

# 加载模型，同时加载了模型的结构、权重等信息
new_model = keras.models.load_model('path_to_my_model.h5')

1.2 将模型保存为SavedModel格式

我们可以将Sequence和Model模型保存为tensorflow标准的SavedModel格式，这是对tensorflow对象标准的序列化格式，非常推荐使用，不同的是，他不是将模型保存为一个单独的文件，而是有几个文件组成。

# 模型保存，注意：仅仅是多了一个save_format的参数而已
# 注意：这里的'path_to_saved_model'不再是模型名称，仅仅是一个文件夹，模型会保存在这个文件夹之下
model.save('path_to_saved_model', save_format='tf')

# 加载模型，通过指定存放模型的文件夹来加载
new_model = keras.models.load_model('path_to_saved_model')

需要注意的是，这种方式依然会保存模型的所有信息，即“网络结构、权重、配置、优化器状态”四个信息，所以可以接着训练。

另外保存完成之后，会在模型存放的路径下面出现以下几个文件以及文件：

----path_to_saved_model

    |-----assets

    |-----variables

           |------variables.data-00000-of-00002

           |------variables.data-00001-of-00002

           |------variables.index

    |-----saved_model.pb

关于每一个文件是什么含义，存放着什么内容，后面再介绍。

1.3 分别保存模型的结构Architecture和模型权重Weights

（1）仅仅保存模型结构Architecture——两个方法

# 方法一：通过model的get_config方法
config = model.get_config()  # 注意这里并没有存储到磁盘哦！！

# 方法二：通过json格式，可以将模型保存为json格式
json_config = model.to_json()

# 将模型结构以json格式保存成磁盘文件
with open('model_config.json', 'w') as json_file:
    json_file.write(json_config)

模型结构的加载，

# 从前面的config对象加载
new_model = keras.Model.from_config(config)

# 从前面的json_config对象加载，这个对象可以是从磁盘读取的就送对象
new_model = keras.models.model_from_json(json_config)

# 如果从磁盘读取，则
with open('model_config.json') as json_file:
    json_config = json_file.read()
new_model = keras.models.model_from_json(json_config)

（2）保存模型的权重——Weights

不保存权重，仅仅是获取模型的权重，然后给模型设置

# 仅仅是获取权重，不保存
weights = model.get_weights()  

# 给模型设置权重
model.set_weights(weights) 

将权重weights保存在磁盘上，用一个 h5 文件来存储。

# 保存权重到磁盘，注意，这里是一个 h5 文件哦！
model.save_weights('path_to_my_weights.h5')

# 新模型从磁盘加载权重
new_model.load_weights('path_to_my_weights.h5')

下面演示一个完整的过程

# 保存模型结构，存储在磁盘上
json_config = model.to_json()
with open('model_config.json', 'w') as json_file:
    json_file.write(json_config)

# 保存模型权重，也存储在磁盘上
model.save_weights('path_to_my_weights.h5')

# 从磁盘加载模型
with open('model_config.json') as json_file:
    json_config = json_file.read()
new_model = keras.models.model_from_json(json_config)

# 从磁盘加载权重
new_model.load_weights('path_to_my_weights.h5')

1.4 仅仅保存模型权重，保存成checkpoint格式

就一句话

# 注意，这里的'path_to_my_tf_checkpoint'是一个文件夹
# 权重checkpoint 会保存在这个文件夹之下
model.save_weights('path_to_my_tf_checkpoint')

特别注意：

model的save_weights()方法既可以包成一个单独的 h5 权重文件，也可以保存成checkpoint的格式，取决于传入的参数是什么，当字符串是一个 “ xxx .h5” 文件的时候，保存成h5文件，当是一个文件夹的时候，默认保存checkpoint。

为了便于区分，我们在保存checkpoint的时候，可以传入关键字参数 save_format = "tf" 

# 为了和保存成 h5 文件进行区分，通过关键字 save_format
model.save_weights('path_to_my_tf_checkpoint', save_format='tf')

比如我的checkpoint保存代码如下所示：

 model.save_weights('./model-weights/checkpoint_weights', save_format='tf')
'''
运行之后会在 ./model-weights文件夹下面产生如下几个文件，其中checkpoint-weights是文件名称

-----model-weights
     |_____checkpoint
     |_____checkpoint_weights.data-00000-of-00002
     |_____checkpoint_weights.data-00001-of-00002
     |_____checkpoint_weights.index
'''
'''
可以看出，下面的三个文件实际上和SavedModel保存的格式的variables文件夹中的是一样的
'''

注意：保存成的checkpoint文件不是一个，而是好几个文件，里面包含的信息有

• 模型的权重
• 优化器optimizer的状态
• 其他的一些配置信息

相较于完整的模型，他只缺少模型的结构Architecture，因此，也是可以用于重建模型，接着往后训练的，关于如何使用，后面再专门讲解checkpoint文件。

 

1.5 小总结

总结之一——模型保存的四种方法,参见上面

• 完整模型  h5 文件
• 完整的SavedModel格式
• 结构和权重单独存储，其中权重可以是  h5 文件，也可以是checkpoint文件格式

总结之二——几个函数的几种不同用法

（1）model.save()有两种用法——取决于参数是  “文件名称”  还是  “文件夹”

model.save('path_to_my_model.h5')   # 保存成完整的 h5 文件

model.save('path_to_saved_model', save_format='tf')  # 保存完整模型为SavedModel格式

（2）keras.models.load_model()加载完整模型的两种方法，取决于参数是  “文件名称”  还是  “文件夹”

new_model = keras.models.load_model('path_to_my_model.h5')   # 文件名称

new_model = keras.models.load_model('path_to_saved_model')   # 文件夹

（3）model.save_weights()保存模型权重的两种方法，取决于参数是  “文件名称”  还是  “文件夹”

model.save_weights('path_to_my_weights.h5')      # 文件名称

model.save_weights('path_to_my_tf_checkpoint')  # 文件夹，保存checkpoint权重

 （4）new_model.load_weights()加载模型权重的两种方法，取决于参数是  “文件名称”  还是  “文件夹”

new_model.load_weights('path_to_my_weights.h5')        # 文件名称

new_model.load_weights('path_to_my_tf_checkpoint')#文件夹名称

其中（1）（2）搭配使用，（3）（4）搭配使用。

 

二、SubClass Model的模型以及权重保存

统计通过继承Model类来实现的模型子类化，在保存模型上面与keras自己实现的方式有一些区别，我们不能够像Sequence和Model那样，完整的一下子将整个模型（结构、权重、配置、优化器状态）保存成一个完整的 h5 文件，但是我们有以下几种方式可以选择。

（1）仅仅保存checkpoint权重信息

my_model.save_weights('path_to_my_weights', save_format='tf')

关于到底什么是checkpoint，后面会介绍。

（2）将整个模型模型保存为SavedModel格式

# 保存整个模型，保存成SavedModel格式
model.save('path_to_my_model',save_format='tf')

# 从SavedModel加载模型
new_model = keras.models.load_model('path_to_my_model')

但是是用这个方法的前提条件是对于自定义的模型，我是用keras提供的内置训练模式，即通过model.fit来进行训练，才能够使用这种方法保存整个模型，但是如果没有使用任何的fit、predict方法，则会提示错误，显示不知道模型的输入形状，需要调用fit或者是predict之后才行，提示我们使用model._set_inputs()来进行设置。

但是有时候我需要定义自己的训练流程，不想使用fit进行训练，那怎么办呢？我们可以这么做，在模型训练之前设置模型的输入形状。

通过model._set_inputs()来设置：

# 创建自定义模型
model = KeypointNetwork()  

# 对于自定义模型，给模型制定一个输入形状，这对于后面模型的保存以及加载，是有必要的
# 通过TensorSpec创建一个“无实际数据的张量”，指定它的形状，作为模型的输入
shape = tf.TensorSpec(shape = (batch_size,210), dtype=tf.dtypes.float32, name=None)

# 设置模型的输入
model._set_inputs(shape)  # 设置模型的输入形状

# 开始训练流程
for epoch in range(1,epochs+1):
    ... ...
    ... ...


# 现在可以保存完整的模型了
model.save('./model-weights/checkpoint_weights',save_format='tf')

关于model._set_inputs()的具体解释，可以自己查看帮助文档

help(tf.keras.Model._set_inputs)

（3）使用tf函数将整个模型保存成SavedModel格式

实际上这个方法和方法（2）实现的功能是一样的，只不过换了一个函数来实现，注意问题也是一样的，如果使用了fit来训练，则可以直接保存，否则需要通过model._set_inputs来指定模型的输入。

# 保存整个模型成SavedModel格式，注意参数'my_saved_model'是文件夹名称，不是文件哦！
tf.saved_model.save(model,'my_saved_model')

# 加载SavedModel格式的模型
restored_saved_model = keras.models.load_model('my_saved_model')

总结：对于子类化的模型SubClass Model，不能够将整个模型保存成 h5 文件，也不能够将权重保存成 h5 文件，只能够将整个模型保存成SavedModel模型，将权重保存成checkpoint格式的权重。