[mnist]: Address comments from PR

https://github.com/tensorflow/models/pull/3024

official/mnist/README.md

@@ -42,7 +42,7 @@ The SavedModel will be saved in a timestamped directory under `/tmp/mnist_saved_
 Use [`saved_model_cli`](https://www.tensorflow.org/programmers_guide/saved_model#cli_to_inspect_and_execute_savedmodel) to inspect and execute the SavedModel.
 
 ```
-saved_model_cli run --dir /tmp/mnist_saved_model/TIMESTAMP --tag_set serve --signature_def classify --inputs image_raw=examples.npy
+saved_model_cli run --dir /tmp/mnist_saved_model/TIMESTAMP --tag_set serve --signature_def classify --inputs image=examples.npy
 ```
 
 `examples.npy` contains the data from `example5.png` and `example3.png` in a numpy array, in that order. The array values are normalized to values between 0 and 1.

official/mnist/mnist.py

@@ -126,10 +126,14 @@ class Model(object):
     return self.fc2(y)
 
 
-def predict_spec(model, image):
-  """EstimatorSpec for predictions."""
+def model_fn(features, labels, mode, params):
+  """The model_fn argument for creating an Estimator."""
+  model = Model(params['data_format'])
+  image = features
   if isinstance(image, dict):
-    image = image['image']
+    image = features['image']
+
+  if mode == tf.estimator.ModeKeys.PREDICT:
     logits = model(image, training=False)
     predictions = {
         'classes': tf.argmax(logits, axis=1),
@@ -141,10 +145,7 @@ def predict_spec(model, image):
         export_outputs={
             'classify': tf.estimator.export.PredictOutput(predictions)
         })
-
-
-def train_spec(model, image, labels):
-  """EstimatorSpec for training."""
+  if mode == tf.estimator.ModeKeys.TRAIN:
     optimizer = tf.train.AdamOptimizer(learning_rate=1e-4)
     logits = model(image, training=True)
     loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
@@ -158,10 +159,7 @@ def train_spec(model, image, labels):
         mode=tf.estimator.ModeKeys.TRAIN,
         loss=loss,
         train_op=optimizer.minimize(loss, tf.train.get_or_create_global_step()))
-
-
-def eval_spec(model, image, labels):
-  """EstimatorSpec for evaluation."""
+  if mode == tf.estimator.ModeKeys.EVAL:
     logits = model(image, training=False)
     loss = tf.losses.softmax_cross_entropy(onehot_labels=labels, logits=logits)
     return tf.estimator.EstimatorSpec(
@@ -175,17 +173,6 @@ def eval_spec(model, image, labels):
         })
 
 
-def model_fn(features, labels, mode, params):
-  """The model_fn argument for creating an Estimator."""
-  model = Model(params['data_format'])
-  if mode == tf.estimator.ModeKeys.PREDICT:
-    return predict_spec(model, features)
-  if mode == tf.estimator.ModeKeys.TRAIN:
-    return train_spec(model, features, labels)
-  if mode == tf.estimator.ModeKeys.EVAL:
-    return eval_spec(model, features, labels)
-
-
 def main(unused_argv):
   data_format = FLAGS.data_format
   if data_format is None:
@@ -227,7 +214,7 @@ def main(unused_argv):
   if FLAGS.export_dir is not None:
     image = tf.placeholder(tf.float32, [None, 28, 28])
     input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
-        'image': tf.placeholder(tf.float32, [None, 28, 28])
+        'image': image,
     })
     mnist_classifier.export_savedmodel(FLAGS.export_dir, input_fn)