как сохранить модель тензорного потока с помощью tf.estimator

Question

У меня есть следующий пример кода для обучения и оценки модели cnn mnist с использованием api инструмента оценки tenorflow:

 def model_fn(features, labels, mode):
        images = tf.reshape(features, [-1, 28, 28, 1])
        model = Model()
        logits = model(images)

        predicted_logit = tf.argmax(input=logits, axis=1, output_type=tf.int32)

        if mode == tf.estimator.ModeKeys.PREDICT:
            probabilities = tf.nn.softmax(logits)

            predictions = {
                'predicted_logit': predicted_logit,
                'probabilities': probabilities
            }
            return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

        else:
            ...

    def mnist_train_and_eval(_):
        train_data, train_labels, eval_data, eval_labels, val_data, val_labels = get_mnist_data()

        # Create a input function to train
        train_input_fn = tf.estimator.inputs.numpy_input_fn(
            x= train_data,
            y=train_labels,
            batch_size=_BATCH_SIZE,
            num_epochs=1,
            shuffle=True)

        # Create a input function to eval
        eval_input_fn = tf.estimator.inputs.numpy_input_fn(
            x= eval_data,
            y=eval_labels,
            batch_size=_BATCH_SIZE,
            num_epochs=1,
            shuffle=False)

        # Create a estimator with model_fn
        image_classifier = tf.estimator.Estimator(model_fn=model_fn, model_dir=_MODEL_DIR)

        # Finally, train and evaluate the model after each epoch
        for _ in range(_NUM_EPOCHS):
            image_classifier.train(input_fn=train_input_fn)
            metrics = image_classifier.evaluate(input_fn=eval_input_fn)

Как я могу использовать estimator.export_savedmodel, чтобы сохранить обученную модель для последующего вывода? Как мне написать serve_input_receiver_fn?

Большое спасибо за помощь!

Answer 1

Вы создаете функцию со словарем входных функций. Заполнитель должен соответствовать форме вашего изображения с первым измерением для batch_size.

def serving_input_receiver_fn():
  x = tf.placeholder(tf.float32, [None, Shape])
  inputs = {'x': x}
  return tf.estimator.export.ServingInputReceiver(features=inputs, receiver_tensors=inputs)

Или вы можете использовать TensorServingInputReceiver, который не требует наложения текста

inputs = tf.placeholder(tf.float32, [None, 32*32*3])
tf.estimator.export.TensorServingInputReceiver(inputs, inputs)

Эта функция возвращает новый экземпляр ServingInputReceiver, который передается в export_savedmodel или tf.estimator.FinalExporter

...
image_classifier.export_savedmodel(saved_dir, serving_input_receiver_fn)