tf2.X中将模型保存为savedmodel

背景

将一个计算欧式距离的小模型转换为savedmodel。

代码

import tensorflow as tf


class EuclideanDisTanceNet(tf.keras.Model):
  def __init__(self):
    super(EuclideanDisTanceNet, self).__init__()
    

  def call(self, input):
    sub = tf.subtract(input[0], input[1])
    print('sub>', sub)
    pow = tf.pow(sub, 2)
    print('pow>', pow)
    rd = tf.reduce_sum(pow, axis=3)
    print('rd>', rd)
    sq = tf.math.sqrt(rd)
    return sq


input1 = tf.constant(3, dtype=tf.float32, shape=[1, 1, 500, 512])   
input2 = tf.constant(1, dtype=tf.float32, shape=[1, 1, 1, 512])

model = EuclideanDisTanceNet()

# out = model([input1, input2])
# print("out>", out)


out = model.predict([input1, input2])
print("out>", out)
model.save("./savedmodle", include_optimizer=False)

特别注意

如果没有调用predict()，则上面的代码会报错：

cannot be saved because the input shapes have not been set. Usually, input shapes are automatically determined from calling `.fit()` or `.predict()`. To manually set the shapes, call `model.build(input_shape)`.

参考资料

在 TensorFlow 2 中使用 TF Hub 中的 SavedModel

tensorflow2.x:构建tf.keras.Model实例的几种方式