Fixed Keras Integ test

Author: mvsusp

tests/data/cifar_10/source/keras_cnn_cifar_10.py

@@ -10,17 +10,12 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
+from __future__ import absolute_import, division, print_function
 
 import tensorflow as tf
-from tensorflow.python.keras.layers import InputLayer, Conv2D, Activation, MaxPooling2D, Dropout, Flatten, Dense
+from tensorflow.python.keras.layers import Activation, Conv2D, Dense, Dropout, Flatten, MaxPooling2D
 from tensorflow.python.keras.models import Sequential
-from tensorflow.python.keras.optimizers import RMSprop
-from tensorflow.python.saved_model.signature_constants import PREDICT_INPUTS
+from tensorflow.python.training.rmsprop import RMSPropOptimizer
 
 HEIGHT = 32
 WIDTH = 32
@@ -29,7 +24,7 @@
 NUM_DATA_BATCHES = 5
 NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 10000 * NUM_DATA_BATCHES
 BATCH_SIZE = 128
-INPUT_TENSOR_NAME = PREDICT_INPUTS
+INPUT_TENSOR_NAME = 'inputs_input'  # needs to match the name of the first layer + "_input"
 
 
 def keras_model_fn(hyperparameters):
@@ -43,10 +38,7 @@ def keras_model_fn(hyperparameters):
     """
     model = Sequential()
 
-    # TensorFlow Serving default prediction input tensor name is PREDICT_INPUTS.
-    # We must conform to this naming scheme.
-    model.add(InputLayer(input_shape=(HEIGHT, WIDTH, DEPTH), name=PREDICT_INPUTS))
-    model.add(Conv2D(32, (3, 3), padding='same'))
+    model.add(Conv2D(32, (3, 3), padding='same', name='inputs', input_shape=(HEIGHT, WIDTH, DEPTH)))
     model.add(Activation('relu'))
     model.add(Conv2D(32, (3, 3)))
     model.add(Activation('relu'))
@@ -67,19 +59,17 @@ def keras_model_fn(hyperparameters):
     model.add(Dense(NUM_CLASSES))
     model.add(Activation('softmax'))
 
-    _model = tf.keras.Model(inputs=model.input, outputs=model.output)
-
-    opt = RMSprop(lr=hyperparameters['learning_rate'], decay=hyperparameters['decay'])
+    opt = RMSPropOptimizer(learning_rate=hyperparameters['learning_rate'], decay=hyperparameters['decay'])
 
-    _model.compile(loss='categorical_crossentropy',
-                   optimizer=opt,
-                   metrics=['accuracy'])
+    model.compile(loss='categorical_crossentropy',
+                  optimizer=opt,
+                  metrics=['accuracy'])
 
-    return _model
+    return model
 
 
 def serving_input_fn(hyperpameters):
-    inputs = {PREDICT_INPUTS: tf.placeholder(tf.float32, [None, 32, 32, 3])}
+    inputs = {INPUT_TENSOR_NAME: tf.placeholder(tf.float32, [None, 32, 32, 3])}
     return tf.estimator.export.ServingInputReceiver(inputs, inputs)
 
 

tests/integ/test_tf_keras.py

@@ -33,7 +33,7 @@ def test_keras(sagemaker_session, tf_full_version):
                                source_dir=script_path,
                                role='SageMakerRole', sagemaker_session=sagemaker_session,
                                hyperparameters={'learning_rate': 1e-4, 'decay': 1e-6},
-                               training_steps=500, evaluation_steps=5,
+                               training_steps=50, evaluation_steps=5,
                                train_instance_count=1, train_instance_type='ml.c4.xlarge',
                                train_max_run=45 * 60)