Add keras notebook example

sagemaker-python-sdk/tensorflow_keras_cifar10/cifar10_cnn.py

@@ -0,0 +1,203 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import tensorflow as tf
+from tensorflow.python.estimator.export.export import build_raw_serving_input_receiver_fn
+from tensorflow.python.keras._impl.keras.engine.topology import InputLayer
+from tensorflow.python.keras._impl.keras.layers import Conv2D, Activation, MaxPooling2D, Dropout, Flatten, Dense
+from tensorflow.python.keras._impl.keras.models import Sequential
+from tensorflow.python.keras._impl.keras.optimizers import rmsprop
+from tensorflow.python.saved_model.signature_constants import PREDICT_INPUTS
+
+HEIGHT = 32
+WIDTH = 32
+DEPTH = 3
+NUM_CLASSES = 10
+NUM_DATA_BATCHES = 5
+NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN = 10000 * NUM_DATA_BATCHES
+BATCH_SIZE = 1
+
+
+def keras_model_fn(hyperparameters):
+    """keras_model_fn receives hyperparameters from the training job and returns a compiled keras model.
+    The model will transformed in a TensorFlow Estimator before training and it will saved in a TensorFlow Serving
+    SavedModel in the end of training.
+
+    Args:
+        hyperparameters: The hyperparameters passed to SageMaker TrainingJob that runs your TensorFlow training
+                         script.
+    Returns: A compiled Keras model
+    """
+    model = Sequential()
+
+    # TensorFlow Serving default prediction input tensor name is PREDICT_INPUTS. I will keep the same name for the
+    # InputLayer
+    model.add(InputLayer(input_shape=(HEIGHT, WIDTH, DEPTH), name=PREDICT_INPUTS))
+    model.add(Conv2D(32, (3, 3), padding='same'))
+    model.add(Activation('relu'))
+    model.add(Conv2D(32, (3, 3)))
+    model.add(Activation('relu'))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.25))
+
+    model.add(Conv2D(64, (3, 3), padding='same'))
+    model.add(Activation('relu'))
+    model.add(Conv2D(64, (3, 3)))
+    model.add(Activation('relu'))
+    model.add(MaxPooling2D(pool_size=(2, 2)))
+    model.add(Dropout(0.25))
+
+    model.add(Flatten())
+    model.add(Dense(512))
+    model.add(Activation('relu'))
+    model.add(Dropout(0.5))
+    model.add(Dense(NUM_CLASSES))
+    model.add(Activation('softmax'))
+
+    opt = rmsprop(lr=0.0001, decay=1e-6)
+
+    model.compile(loss='categorical_crossentropy',
+                  optimizer=opt,
+                  metrics=['accuracy'])
+
+    print(model.summary())
+
+    return model
+
+
+def serving_input_fn(hyperparameters):
+    """This function defines the placeholders that will be added to the model during serving.
+    The function returns a tf.estimator.export.ServingInputReceiver object, which packages the placeholders and the
+    resulting feature Tensors together.
+
+    For more information: https://github.com/aws/sagemaker-python-sdk#creating-a-serving_input_fn
+
+    Args:
+        hyperparameters: The hyperparameters passed to SageMaker TrainingJob that runs your TensorFlow training
+                        script.
+    Returns: ServingInputReceiver or fn that returns a ServingInputReceiver
+    """
+
+    # Notice that the input placeholder has the same input shape as the Keras model input
+    tensor = tf.placeholder(tf.float32, shape=[None, HEIGHT, WIDTH, DEPTH])
+
+    # the features key PREDICT_INPUTS matches the Keras Input Layer name
+    features = {PREDICT_INPUTS: tensor}
+    return build_raw_serving_input_receiver_fn(features)
+
+
+def train_input_fn(training_dir, hyperparameters):
+    return _input(tf.estimator.ModeKeys.TRAIN,
+                  batch_size=BATCH_SIZE, data_dir=training_dir)
+
+
+def eval_input_fn(training_dir, hyperparameters):
+    return _input(tf.estimator.ModeKeys.EVAL,
+                  batch_size=BATCH_SIZE, data_dir=training_dir)
+
+
+def _input(mode, batch_size, data_dir):
+    """Input_fn using the contrib.data input pipeline for CIFAR-10 dataset.
+
+  Args:
+    mode: Standard names for model modes (tf.estimators.ModeKeys).
+    batch_size: The number of samples per batch of input requested.
+  """
+    dataset = _record_dataset(_filenames(mode, data_dir))
+
+    # For training repeat forever.
+    if mode == tf.estimator.ModeKeys.TRAIN:
+        dataset = dataset.repeat()
+
+    dataset = dataset.map(_dataset_parser, num_threads=1,
+                          output_buffer_size=2 * batch_size)
+
+    # For training, preprocess the image and shuffle.
+    if mode == tf.estimator.ModeKeys.TRAIN:
+        dataset = dataset.map(_train_preprocess_fn, num_threads=1,
+                              output_buffer_size=2 * batch_size)
+
+        # Ensure that the capacity is sufficiently large to provide good random
+        # shuffling.
+        buffer_size = int(NUM_EXAMPLES_PER_EPOCH_FOR_TRAIN * 0.4) + 3 * batch_size
+        dataset = dataset.shuffle(buffer_size=buffer_size)
+
+    # Subtract off the mean and divide by the variance of the pixels.
+    dataset = dataset.map(
+        lambda image, label: (tf.image.per_image_standardization(image), label),
+        num_threads=1,
+        output_buffer_size=2 * batch_size)
+
+    # Batch results by up to batch_size, and then fetch the tuple from the
+    # iterator.
+    iterator = dataset.batch(batch_size).make_one_shot_iterator()
+    images, labels = iterator.get_next()
+
+    return {PREDICT_INPUTS: images}, labels
+
+
+def _train_preprocess_fn(image, label):
+    """Preprocess a single training image of layout [height, width, depth]."""
+    # Resize the image to add four extra pixels on each side.
+    image = tf.image.resize_image_with_crop_or_pad(image, HEIGHT + 8, WIDTH + 8)
+
+    # Randomly crop a [HEIGHT, WIDTH] section of the image.
+    image = tf.random_crop(image, [HEIGHT, WIDTH, DEPTH])
+
+    # Randomly flip the image horizontally.
+    image = tf.image.random_flip_left_right(image)
+
+    return image, label
+
+
+def _dataset_parser(value):
+    """Parse a CIFAR-10 record from value."""
+    # Every record consists of a label followed by the image, with a fixed number
+    # of bytes for each.
+    label_bytes = 1
+    image_bytes = HEIGHT * WIDTH * DEPTH
+    record_bytes = label_bytes + image_bytes
+
+    # Convert from a string to a vector of uint8 that is record_bytes long.
+    raw_record = tf.decode_raw(value, tf.uint8)
+
+    # The first byte represents the label, which we convert from uint8 to int32.
+    label = tf.cast(raw_record[0], tf.int32)
+
+    # The remaining bytes after the label represent the image, which we reshape
+    # from [depth * height * width] to [depth, height, width].
+    depth_major = tf.reshape(raw_record[label_bytes:record_bytes],
+                             [DEPTH, HEIGHT, WIDTH])
+
+    # Convert from [depth, height, width] to [height, width, depth], and cast as
+    # float32.
+    image = tf.cast(tf.transpose(depth_major, [1, 2, 0]), tf.float32)
+
+    return image, tf.one_hot(label, NUM_CLASSES)
+
+
+def _record_dataset(filenames):
+    """Returns an input pipeline Dataset from `filenames`."""
+    record_bytes = HEIGHT * WIDTH * DEPTH + 1
+    return tf.contrib.data.FixedLengthRecordDataset(filenames, record_bytes)
+
+
+def _filenames(mode, data_dir):
+    """Returns a list of filenames based on 'mode'."""
+    data_dir = os.path.join(data_dir, 'cifar-10-batches-bin')
+
+    assert os.path.exists(data_dir), ('Run cifar10_download_and_extract.py first '
+                                      'to download and extract the CIFAR-10 data.')
+
+    if mode == tf.estimator.ModeKeys.TRAIN:
+        return [
+            os.path.join(data_dir, 'data_batch_%d.bin' % i)
+            for i in range(1, NUM_DATA_BATCHES + 1)
+        ]
+    elif mode == tf.estimator.ModeKeys.EVAL:
+        return [os.path.join(data_dir, 'test_batch.bin')]
+    else:
+        raise ValueError('Invalid mode: %s' % mode)