Tensorflow 2.0 如何从 Keras 模型导出预测和训练模式

Question

I have a model sub-classed from tf.keras.model . I have write a call -and predict methods. when I export the model it seems that only the output from call method is serialized. The simple code to illustrate the problem is below

class SimpleModel(tf.keras.Model):

def __init__(self):
    super(SimpleModel, self).__init__()

    self.layer1 = keras.layers.Flatten(input_shape=(28, 28))
    self.layer2 = keras.layers.Dense(128, activation='relu')
    self.dropout = keras.layers.Dropout(0.5)
    self.layer3 = keras.layers.Dense(10, activation='softmax')

def call(self, x, training=False):

    x = self.layer1(x)
    x = self.layer2(x)
    if training:
        x = self.dropout(x)
    return self.layer3(x)

def predict(self, x):

    return tf.argmax(self(x))

fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

model = SimpleModel()

model.compile(optimizer='adam',
          loss='sparse_categorical_crossentropy',
          metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=10)

model.save('tf_test', save_format='tf')

When I inspect the saved model using the saved-model-cli the output is as below

The given SavedModel SignatureDef contains the following input(s):
  inputs['input_1'] tensor_info:
      dtype: DT_UINT8
      shape: (-1, 28, 28)
      name: serving_default_input_1:0
The given SavedModel SignatureDef contains the following output(s):
  outputs['output_1'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 10)
      name: StatefulPartitionedCall:0
Method name is: tensorflow/serving/predict

The output is exactly same without the predict method. So how can I get the prediction tensor included in the serialization and does the training mode need to be serialized. In tensorflow 1.x I can simply save the prediction tensor and training mode tensor using the following line of code

tf.saved_model.simple_save(sess, 'tf_test', inputs={"x": x, "mode": training_mode}, outputs={"predictions": predictions})

Answer 1

I think I find the solution to the export problem from the tensorflow documentation Using save model format section Exporting custom models . The solution is to use tf.Module instead tf.keras.Model and use the tf.function annotation at the top of the function one wants to import. The working code is illustrated below

class SimpleModel(tf.Module):

def __init__(self):
    super(SimpleModel, self).__init__()

    self.layer1 = keras.layers.Flatten(input_shape=(28, 28))
    self.layer2 = keras.layers.Dense(128, activation='relu')
    self.dropout = keras.layers.Dropout(0.5)
    self.layer3 = keras.layers.Dense(10, activation='softmax')

def __call__(self, x, training=False):

    x = self.layer1(x)
    x = self.layer2(x)
    # if training:
    #     x = self.dropout(x)
    x = self.layer3(x)
    return x

@tf.function(input_signature=[tf.TensorSpec([None, 28, 28], tf.float32)])
def predict(self, x):

    return tf.argmax(self(x), axis=1)


def loss(m, x, y):
    logits = m(x, True)
    return tf.keras.losses.categorical_crossentropy(tf.reshape(tf.one_hot(y, 10), (y.size, 10)), logits)


def grad(m, x, y):
    with tf.GradientTape() as tape:
    lv = loss(m, x, y)
    return lv, tape.gradient(lv, m.trainable_variables)


fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = 
fashion_mnist.load_data()

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
           'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

model = SimpleModel()

# model.compile(optimizer='adam',
#               loss='sparse_categorical_crossentropy',
#               metrics=['accuracy'])
#
# model.fit(train_images, train_labels, epochs=10)

epochs = 100
batch_size = 1000
optimizer = tf.keras.optimizers.SGD(learning_rate=1e-5)

for epoch in range(epochs):

    costs = []
    for k in range(int(train_images.shape[0] / batch_size)):
        start_index = k * batch_size
        end_index = (k + 1) * batch_size
        batch_x, batch_y = train_images[start_index:end_index, :, :], 
        train_labels[start_index:end_index]

        loss_value, grads = grad(model, batch_x, np.reshape(batch_y, (1000, 1)))
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        costs.append(loss_value.numpy())
    print("epoch %d of %d, cost %f" % (epoch, 10, np.mean(costs)))

signatures = {"serving_default": model.predict}

tf.saved_model.save(model, 'tf_test', signatures)

No the output of the saved_model_cli is

The given SavedModel SignatureDef contains the following input(s):
    inputs['x'] tensor_info:
    dtype: DT_FLOAT
    shape: (-1, 28, 28)
    name: serving_default_x:0
The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
    dtype: DT_INT64
    shape: (-1)
    name: StatefulPartitionedCall:0
Method name is: tensorflow/serving/predict

which corresponds to the output of predict method outputting single label