背景
将一个计算欧式距离的小模型转换为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)`.