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

Question

我有一个从 tf.keras.model 子类化的模型。我写了一个调用和预测方法。当我导出模型时，似乎只有调用方法的输出被序列化。下面是说明问题的简单代码

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')

当我使用保存的模型-cli 检查保存的模型时，输出如下

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

没有使用预测方法时，输出完全相同。那么如何获得序列化中包含的预测张量以及训练模式是否需要序列化。在tensorflow 1.x中，我可以使用以下代码行简单地保存预测张量和训练模式张量

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

Answer 1

我想我从tensorflow文档Using save model format中找到了导出问题的解决方案导出自定义模型部分。解决方案是使用 tf.Module 而不是 tf.keras.Model 并在要导入的函数顶部使用 tf.function 注释。工作代码如下所示

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)

否，saved_model_cli 的输出是

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

对应于预测方法输出单个标签的输出