I get an error when I load a tensorflow2.0 model - tensorflow2.0

Question

I am learning a simple model to perform a linear regression and then I save the model
class NN(tf.keras.Model):
def __init__(self):
super(NN, self).__init__()
L = 20
self.W1 = tf.Variable(tf.random.truncated_normal([1, L], stddev=math.sqrt(3)))
self.B1 = tf.Variable(tf.random.truncated_normal([1, L], stddev=1.0))
self.W2 = tf.Variable(tf.random.truncated_normal([L, 1], stddev=math.sqrt(3/L)))
self.B2 = tf.Variable(tf.zeros([1]))
def call(self, inputs):
Z1 = tf.matmul(inputs, self.W1) + self.B1
Y1 = tf.nn.tanh(Z1)
Y = tf.matmul(Y1, self.W2) + self.B2
return Y
# The loss function to be optimized
def loss(model, X, Y_):
error = model(X) - Y_
return tf.reduce_mean(tf.square(error))
model = NN()
optimizer = tf.optimizers.Adam(learning_rate=0.001)
bsize = 20
# You can call this function in a loop to train the model, bsize samples at a time
def training_step(i):
# read data
x_batch, y_batch = func.next_batch(bsize)
x_batch = np.reshape(x_batch, (bsize,1))
y_batch = np.reshape(y_batch, (bsize,1))
# compute training values
loss_fn = lambda: loss(model, x_batch, y_batch)
optimizer.minimize(loss_fn, [model.W1, model.B1, model.W2, model.B2])
if i%5000 == 0:
l = loss(model, x_batch, y_batch)
print(str(i) + ": epoch: " + str(func._epochs_completed) + ": loss: " + str(l.numpy()))
for i in range(50001):
training_step(i)
# save the model
tf.saved_model.save(model, "my_file")
and then I am trying to load the model with the following lines following tensorflow documentation:
model = tf.saved_model.load("my_file")
f = model.signatures["serving_default"]
y = f(x)
However I get the following error message:
f = model.signatures["serving_default"]
File "my_file/signature_serialization.py", line 195, in __getitem__
return self._signatures[key]
KeyError: 'serving_default'
What is wrong ? Why serving_default is not defined ?

Answer 1

I solved the problem by adding a third argument to the tf.saved_model.save function
tf.saved_model.save(model, "myfile", signatures=model.call.get_concrete_function(tf.TensorSpec(shape=[None,1], dtype=tf.float32, name="inp")))
and by adding the #tf.function above the call method
class NN(tf.keras.Model):
def __init__(self):
super(NN, self).__init__()
L = 20
self.W1 = tf.Variable(tf.random.truncated_normal([1, L], stddev=math.sqrt(3)))
self.B1 = tf.Variable(tf.random.truncated_normal([1, L], stddev=1.0))
self.W2 = tf.Variable(tf.random.truncated_normal([L, 1], stddev=math.sqrt(3/L)))
self.B2 = tf.Variable(tf.zeros([1]))
#tf.function
def call(self, X):
Z1 = tf.matmul(X, self.W1) + self.B1
Y1 = tf.nn.tanh(Z1)
Y = tf.matmul(Y1, self.W2) + self.B2
return Y