Question

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

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

当我使用save-model-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)

save_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

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

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

1 个答案: