基于tf2.x的keras训练的模型转成tf1.x版本
写在前面
我训练使用的是TF2.3，部分业务场景需要将模型转到TF1.14版本进行部署。
训练的模型是基于keras的预训练，外加自己添了一点东西。
TF2.x的模型转到TF1.x之后存在精度损失，因模型而异。
第一步：保证在tf1.x环境下正确导入tf2.x的权重
在TF2.x只保存模型的权重
只保存模型的权重，能够保证在TF1.x环境下，通过定义模型结构可以正确的导入权重，不会引入TF2.x特有的OP或者属性，避免一些莫名奇妙的错误。
tf2_model = load_trained_keras_model()
tf2_model.trainable = False
tf2_model.save("weights.hdf5")
在TF1.x下定义模型结构并导入TF2.x保存的权重
在TF2.x的keras中有些Layer不存在于TF1.x中，比如keras.layers.Normalizaiton，在定义模型结构时，可以对一些不存在的op进行自定义生成，保证Variable的名字一致，这是因为模型导入权重有两种形式：（1）根据网络的拓扑逻辑自动填入权重；（2）根据变量名称填入权重。
在定义网络结构时，一定要和原来保持一致，例如：
inputs = tf.keras.Input(shape=(224, 224, 3))
pretrain = tf.keras.applications.EfficientNetB0(include_top=False, input_tensor=inputs, pooling=None)
x = pretrain(inputs)
x = tf.keras.layers.Conv2D(512, kernel_size=1, use_bias=False)(x)
outputs = tf.keras.layers.GlobalAvgPool2D()(x)
model = tf.keras.Model(inputs, outputs, name="my_model")
model.summary()
可以看到整个模型只有寥寥几层，预训练模型整个被封装成一层，我们在自定义模型结构时，
也要保证efficientB0整体作为一层，这样才能保证在TF1.x下能够按照网络的拓扑逻辑
或者根据变量名进行权重的恢复，如果无法正常恢复，模型会自动初始化，要记得验证一下权重。
在TF1.x下可以这样操作：
from net import EfficientNetB0
inputs = tf.keras.Input(shape=(224, 224, 3))
cust_effb0 = EfficientNetB0(inputs)
model = tf.keras.Model(inputs, cust_effb0, name="efficientnetb0")
x = model (inputs)
x = tf.keras.layers.Conv2D(512, kernel_size=1, use_bias=False)(x)
outputs = tf.keras.layers.GlobalAvgPool2D()(x)
model = tf.keras.Model(inputs, outputs, name="my_model")
model.trainable=False
model.load_weights("weights.hdf5")
第二步，转成TF1.x的saved_model
通过第一步在TF1.x的环境下正确导入权重后，我尝试去通过model.save以及model.save_weights去保存模型或者权重，但都报同一个错误（can not pickle ‘_thread.RLOCK.xxxxx’），大概这个，没截图，记不清楚了。百思不得其解，错误原因是啥，只好尝试转成saved_model格式，即pb格式。
saved_model_dir = "1/"
sess = tf.compat.v1.keras.backend.get_session()
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(saved_model_dir)
input_tensor = {"input_image": tf.compat.v1.saved_model.build_tensor_info(model.input)}
output_tensor = {"output_res": tf.compat.v1.saved_model.build_tensor_info(model.output)}
signature_def = tf.compat.v1.saved_model.build_signature_def(inputs=input_tensor,
outputs=output_tensor,
method_name=tf.saved_model.PREDICT_METHOD_NAME)
builder.add_meta_graph_and_variables(sess,
[tf.compat.v1.saved_model.SERVING],
signature_def_map={"pred": signature_def},
strip_default_attrs=True
)
builder.save()
到此，基本结束，可参考此篇文章使用tensorflow-serving进行部署。
还可以假如一些后处理操作，再次进行保存。
with tf.compat.v1.Graph().as_default() as g:
with tf.compat.v1.Session() as sess:
saved_path = f"1/"
meta_graph_def = tf.compat.v1.saved_model.load(sess, ["serve"], saved_path)
signature = meta_graph_def.signature_def
input_tensor_name = signature["pred_multi"].inputs["input_image"].name
output_tensor_name = signature["pred_multi"].outputs["output_res"].name
inputs = g.get_tensor_by_name(input_tensor_name)
logits = g.get_tensor_by_name(output_tensor_name)
x = tf.math.softmax(logits)
outputs = tf.argmax(x, axis=-1)
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(f"2/")
input_tensor = {"input_image": tf.compat.v1.saved_model.build_tensor_info(inputs)}
output_tensor = {"output_res": tf.compat.v1.saved_model.build_tensor_info(outputs)}
signature_def = tf.compat.v1.saved_model.build_signature_def(inputs=input_tensor,
outputs=output_tensor,
method_name=tf.saved_model.PREDICT_METHOD_NAME)
builder.add_meta_graph_and_variables(sess,
[tf.compat.v1.saved_model.SERVING],
signature_def_map={"pred_multi": signature_def},
strip_default_attrs=True
)
builder.save()
|