TensorFlow2模型保存与加载 自定义预训练模型加载

tensorflow模型的加载机制

模型的保存与恢复

参考指南

计算图与权重保存与恢复

将模型的架构与参数都保存下来，好处是在加载模型时无需重新构建网络，直接将网络加载进行就能够进行预测，因此已经包括了计算图。

SavedModel格式

模型保存：model.save或tf.keras.models.save_model
加载加载：tf.keras.models.load_model

h5格式

函数与SavedModel一样，不同的是模型的名称以 .h5 为后缀。

1. 通过 model.add_loss() 和 model.add_metric() 添加的外部损失和指标不会被保存
2. 已保存的文件中不包含自定义对象（如自定义层）的计算图

权重保存与恢复

参考指南

内存中的权重迁移

取权重：tf.keras.layers.Layer.get_weights()
设置权重：tf.keras.layers.Layer.set_weights()

磁盘中权重迁移

1. TensorFlow 检查点
   模型保存mdoel.save_weights()、tf.train.Checkpoint().save()
   模型加载model.load_weight()、tf.train.Checkpoint().restore()
2. h5
   同样的函数，模型名称后缀为 .h

预训练模型的加载

基于tf.train.Checkpoint的加载方式(成功)

本加载方式主要以分类网络作为示例，预训练模型为分类网络的主干网络，网络的类别数量改变，因此仅加载主干网络参数。
tf.train.Checkpoint函数函数能够输入自定义键值，实际操作其实是在保存时在变量名称的前面加上自定义键值进行限制，因此在加载时候只有相同键值才能加载进去，由此来进行决定加载哪一层的参数。

保存模型

model = create_model()
cp_net = tf.train.Checkpoint(backbone_layer=model.layers[:-1], dense_layer=model.layers[-1])
cp_net.save(filepath)

密集连接层仅有一层，故层数以-1为限，以实际密集连接层数有关

加载模型

to_restore_net = tf.train.Checkpoint(backbone_layer=model.layers[:-1])
to_restore_net.restore(filepath)

基于tf.train.init_from_checkpoint的方式(设想还未实现)

保存模型

model = create_model()
model.save_weights(filepath)

加载模型

预训练模型变量相关信息抽取

预训练参数加载cp_model = tf.train.load_checkpoint(filepath)
获取字典{变量名称：变量形状}cp_model.get_variable_to_shape_map()
获取字典{变量名称：变量类型}cp_model.get_variable_to_dtype_map()
获取var_name的参数cp_model.get_tensor(var_name)

模型变量相关信息抽取

model = create_model()
for layer in model.layers:
	print("layer: ", layer,  "\n",
		   "layer configure: ", layer.get_config())

预训练模型中参数的名称、形状、权重都能获得，搭建模型中参数的名称、形状都能够从get_config()中获取，变量名称与形状相同时可以选择1）建立映射字典，字典作为参数传到tf.train.init_from_checkpoint中 2）利用layer.set_weight进行权重赋予。