TF2模型保存加载

发表于 更新于 阅读次数：

在TF2中保存加载模型，有多种方式。下面介绍下，最主要的几种方式，以及在模型中有Custom Layer时如何保存和加载模型。

由于Google官方在有Custom Layer层 的模型 保存/加载 介绍的非常不全面。这里统一详细介绍下。

.h5 保存/加载

据Google那边的消息，好像他们准备力推TF save_model形式。

一般情况

在利用tf.keras创建的model中直接运用model.save 保存模型的结构和权重为.h5文件。加载时用tf.keras.models.load进行加载。

保存

1
2
3
model = tf.keras.models.Sequential([...])
model.save("model.h5")

加载

1
tf.keras.models.load_model("model.h5")

含有custom layer情况

自定义层：

在编写自定义层时

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class MyDense(tf.keras.layers.Layer):
    # init 可以进行所有与输入无关的初始化
    def __init__(self, units):
        super(MyDense, self).__init__()
        self.units = units
    # build 可以知道输入张量的形状，并可以进行其余的初始化
    def build(self, input_shape):
        self.w1 = self.add_weight(name="w1",shape=[input_shape[-1], self.units])
        self.b1 = self.add_weight(name="b1",shape=(self.units),initializer=tf.keras.initializers.zeros())
    # 在这里进行正向计算
    @tf.function
    def call(self, inputs):
        # return tf.matmul(inputs, self.w1)
        return inputs @ self.w1 + self.b1
    
    # 自定义层配置 在model.save时有用
    def get_config(self):
        config = {
            'units': self.units,
        }
        base_config = super(MyDense,self).get_config()
        return dict(list(config.items()))

在custom layer中的 call() 上添加 @tf.function 可以将前向传播过程中不属于Graph的部分 转化为Graph。

报错一

如果在自定义层中没有重载get_config时，使用model.save() 将会提示如下错误。

NotImplementedError: Layers with arguments in __init__ must override get_config.

如果自定义层中没有重载get_config， 在下文中使用tf.saved_model 中可以正常使用。

因此get_config 只对model.save() 起配置作用。

在custom layer中的 call() 上添加 @tf.function 可以将前向传播过程 转化为Graph。

报错二

若在custom layer的biuld 中创建初始矩阵时，name属性没写，会导致model.save报错

报错提示名字name already exists已存在

File “h5py/_objects.pyx”, line 54, in h5py._objects.with_phil.wrapper
File “h5py/_objects.pyx”, line 55, in h5py._objects.with_phil.wrapper
File “h5py/h5o.pyx”, line 202, in h5py.h5o.link
RuntimeError: Unable to create link (name already exists)

因此正确的写法是 对每一个 矩阵 加上name属性

1
2
self.w1 = self.add_weight(name="w1")
self.b1 = self.add_weight(name="b1")

报错提示名字name already exists已存在

File “h5py/_objects.pyx”, line 54, in h5py._objects.with_phil.wrapper
File “h5py/_objects.pyx”, line 55, in h5py._objects.with_phil.wrapper
File “h5py/h5o.pyx”, line 202, in h5py.h5o.link
RuntimeError: Unable to create link (name already exists)

保存

在含有custom layer时 保存方式和一般情况相同直接，通过model.save(‘model.h5’)即可

加载

在含有custom layer时，加载保存的。.h5 模型文件时需要声明 自定义的层和自定义的损失函数。

在tf.keras.models.load_model中通custom_objects 来指定自定义层，自定义损失函数。

1
model = tf.keras.models.load_model("./model0.h5",custom_objects={"MyDense":MyDense})

TF save_model保存/加载

通过tf.saved_model 保存注意的事情没有tf.keras.models 那么多。但在加载时比较复杂。

特别是，在有自定义层时，tf.save_model 方法在保存时 加载时更加方便，不需要额外的声明自定义层。但，需要在自定义层call 加上**@tf.fuction** 进行修饰。

保存

1
tf.saved_model.save(model, "./test_model_out")

或者直接通过keras model.save 进行保存。 若没有 .h5 后缀，则会自动保存为TF save_model形式。若有.h5 则会保存为.h5文件

1
model.save("./test_model_out")

加载

1
2
3
4
load = tf.save_model.load("./test_model_out")
# 两种方法效果一样
load = tf.keras.models.load_model("./test_model_out")
print(list(loaded.signatures.keys()))

[‘serving_default’]

1
2
infer = loaded.signatures["serving_default"]
print(infer.structured_outputs)

{‘output’: TensorSpec(shape=(None, 1), dtype=tf.float32, name=’output’)}

然后可以直接通过infer进行预测。

输入数据为Tensor变量。

1
2
3
4
5
6
7
test_input_ids   = np.load('./data/test/input_ids.npy')[:10]
test_input_masks = np.load('./data/test/input_masks.npy')[:10]
test_segment_ids = np.load('./data/test/segment_ids.npy')[:10]
# 转换为Tensor变量 其中dtype 需要对应，根据报错信息修改即可 一般为tf.int32 tf.float32
test_input_word_ids =tf.convert_to_tensor(test_input_ids, dtype=tf.int32)
test_input_mask     =tf.convert_to_tensor(test_input_masks, dtype=tf.int32)
test_input_type_ids =tf.convert_to_tensor(test_segment_ids, dtype=tf.int32)
1
2
pre = infer(input_mask=test_input_mask,input_type_ids=test_input_type_ids,input_word_ids=test_input_word_ids)
print(pre)

{‘output’: <tf.Tensor: id=22859, shape=(10, 1), dtype=float32, numpy=
array([[0.16463354],
[0.17504019],
[0.16223168],
[0.18234646],
[0.17594947],
[0.16918331],
[0.15553711],
[0.14790532],
[0.15857093],
[0.17600629]], dtype=float32)>}

通过上面的打印可以看见 layer name = output，可以通过该索引输出。

1
print(pre['output'].numpy())

[[0.16463354]
[0.17504019]
[0.16223168]
[0.18234646]
[0.17594947]
[0.16918331]
[0.15553711]
[0.14790532]
[0.15857093]
[0.17600629]]

总结

在和Google大佬的交谈中发现TF团队准备力推TF save_model 的形式。

https://github.com/tensorflow/models/issues/7643

简单总结下：

模型保存 

1
model.save('file_path')

模型加载

infer()中的数据 需要是Tensor 形式，可以用tf.convert_to_tensor 进行转换

1
2
3
4
5
6
7
load = tf.save_model.load('file_path')
infer = load.signatures['serving_default']
#进行预测
#单输入情况
infer(test_x)											
#多输入情况
infer(x1 =test_x1, x2 = test_x2 )