# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# 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.
"""Convolutional Neural Network Estimator for MNIST, built with tf.layers."""
from __future__ import absolute_import, division, print_function

import argparse
import json
import os

import numpy as np
import tensorflow as tf


def cnn_model_fn(features, labels, mode):
    """Model function for CNN."""
    # Input Layer
    # Reshape X to 4-D tensor: [batch_size, width, height, channels]
    # MNIST images are 28x28 pixels, and have one color channel
    input_layer = tf.reshape(features['x'], [-1, 28, 28, 1])

    # Convolutional Layer #1
    # Computes 32 features using a 5x5 filter with ReLU activation.
    # Padding is added to preserve width and height.
    # Input Tensor Shape: [batch_size, 28, 28, 1]
    # Output Tensor Shape: [batch_size, 28, 28, 32]
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=32,
        kernel_size=[5, 5],
        padding='same',
        activation=tf.nn.relu
    )

    # Pooling Layer #1
    # First max pooling layer with a 2x2 filter and stride of 2
    # Input Tensor Shape: [batch_size, 28, 28, 32]
    # Output Tensor Shape: [batch_size, 14, 14, 32]
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)

    # Convolutional Layer #2
    # Computes 64 features using a 5x5 filter.
    # Padding is added to preserve width and height.
    # Input Tensor Shape: [batch_size, 14, 14, 32]
    # Output Tensor Shape: [batch_size, 14, 14, 64]
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,
        kernel_size=[5, 5],
        padding='same',
        activation=tf.nn.relu
    )

    # Pooling Layer #2
    # Second max pooling layer with a 2x2 filter and stride of 2
    # Input Tensor Shape: [batch_size, 14, 14, 64]
    # Output Tensor Shape: [batch_size, 7, 7, 64]
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)

    # Flatten tensor into a batch of vectors
    # Input Tensor Shape: [batch_size, 7, 7, 64]
    # Output Tensor Shape: [batch_size, 7 * 7 * 64]
    pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])

    # Dense Layer
    # Densely connected layer with 1024 neurons
    # Input Tensor Shape: [batch_size, 7 * 7 * 64]
    # Output Tensor Shape: [batch_size, 1024]
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)

    # Add dropout operation; 0.6 probability that element will be kept
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)

    # Logits layer
    # Input Tensor Shape: [batch_size, 1024]
    # Output Tensor Shape: [batch_size, 10]
    logits = tf.layers.dense(inputs=dropout, units=10)

    predictions = {
        # Generate predictions (for PREDICT and EVAL mode)
        'classes': tf.argmax(input=logits, axis=1),
        # Add `softmax_tensor` to the graph. It is used for PREDICT and by the
        # `logging_hook`.
        'probabilities': tf.nn.softmax(logits, name='softmax_tensor')
    }
    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    # Calculate Loss (for both TRAIN and EVAL modes)
    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    # Configure the Training Op (for TRAIN mode)
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

    # Add evaluation metrics (for EVAL mode)
    eval_metric_ops = {
        'accuracy': tf.metrics.accuracy(
            labels=labels, predictions=predictions['classes'])}
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)


def _load_training_data(base_dir):
    x_train = np.load(os.path.join(base_dir, 'train_data.npy'))
    y_train = np.load(os.path.join(base_dir, 'train_labels.npy'))
    return x_train, y_train


def _load_testing_data(base_dir):
    x_test = np.load(os.path.join(base_dir, 'eval_data.npy'))
    y_test = np.load(os.path.join(base_dir, 'eval_labels.npy'))
    return x_test, y_test


def _parse_args():
    parser = argparse.ArgumentParser()

    # Data, model, and output directories.
    # model_dir is always passed in from SageMaker.
    # By default this is a S3 path under the default bucket.
    parser.add_argument('--model_dir', type=str)
    parser.add_argument('--sm-model-dir', type=str, default=os.environ.get('SM_MODEL_DIR'))
    parser.add_argument('--train', type=str, default=os.environ.get('SM_CHANNEL_TRAINING'))
    parser.add_argument('--hosts', type=list, default=json.loads(os.environ.get('SM_HOSTS')))
    parser.add_argument('--current-host', type=str, default=os.environ.get('SM_CURRENT_HOST'))

    return parser.parse_known_args()


def serving_input_fn():
    inputs = {'x': tf.placeholder(tf.float32, [None, 784])}
    return tf.estimator.export.ServingInputReceiver(inputs, inputs)


if __name__ == '__main__':
    args, _ = _parse_args()

    train_data, train_labels = _load_training_data(args.train)
    eval_data, eval_labels = _load_testing_data(args.train)

    # Create the Estimator
    mnist_classifier = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=args.model_dir)

    # Set up logging for predictions
    # Log the values in the 'Softmax' tensor with label 'probabilities'
    tensors_to_log = {'probabilities': 'softmax_tensor'}
    logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=50)

    # Train the model
    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={'x': train_data},
        y=train_labels,
        batch_size=100,
        num_epochs=None,
        shuffle=True
    )

    # Evaluate the model and print results
    eval_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={'x': eval_data},
        y=eval_labels,
        num_epochs=1,
        shuffle=False
    )

    train_spec = tf.estimator.TrainSpec(train_input_fn, max_steps=20000)
    eval_spec = tf.estimator.EvalSpec(eval_input_fn)
    tf.estimator.train_and_evaluate(mnist_classifier, train_spec, eval_spec)

    if args.current_host == args.hosts[0]:
        mnist_classifier.export_savedmodel(args.sm_model_dir, serving_input_fn)

[((1038, 1080), 'tensorflow.reshape', 'tf.reshape', (["features['x']", '[-1, 28, 28, 1]'], {}), "(features['x'], [-1, 28, 28, 1])\n", (1048, 1080), True, 'import tensorflow as tf\n'), ((1346, 1457), 'tensorflow.layers.conv2d', 'tf.layers.conv2d', ([], {'inputs': 'input_layer', 'filters': '(32)', 'kernel_size': '[5, 5]', 'padding': '"""same"""', 'activation': 'tf.nn.relu'}), "(inputs=input_layer, filters=32, kernel_size=[5, 5],\n    padding='same', activation=tf.nn.relu)\n", (1362, 1457), True, 'import tensorflow as tf\n'), ((1703, 1769), 'tensorflow.layers.max_pooling2d', 'tf.layers.max_pooling2d', ([], {'inputs': 'conv1', 'pool_size': '[2, 2]', 'strides': '(2)'}), '(inputs=conv1, pool_size=[2, 2], strides=2)\n', (1726, 1769), True, 'import tensorflow as tf\n'), ((2015, 2121), 'tensorflow.layers.conv2d', 'tf.layers.conv2d', ([], {'inputs': 'pool1', 'filters': '(64)', 'kernel_size': '[5, 5]', 'padding': '"""same"""', 'activation': 'tf.nn.relu'}), "(inputs=pool1, filters=64, kernel_size=[5, 5], padding=\n    'same', activation=tf.nn.relu)\n", (2031, 2121), True, 'import tensorflow as tf\n'), ((2365, 2431), 'tensorflow.layers.max_pooling2d', 'tf.layers.max_pooling2d', ([], {'inputs': 'conv2', 'pool_size': '[2, 2]', 'strides': '(2)'}), '(inputs=conv2, pool_size=[2, 2], strides=2)\n', (2388, 2431), True, 'import tensorflow as tf\n'), ((2596, 2631), 'tensorflow.reshape', 'tf.reshape', (['pool2', '[-1, 7 * 7 * 64]'], {}), '(pool2, [-1, 7 * 7 * 64])\n', (2606, 2631), True, 'import tensorflow as tf\n'), ((2808, 2877), 'tensorflow.layers.dense', 'tf.layers.dense', ([], {'inputs': 'pool2_flat', 'units': '(1024)', 'activation': 'tf.nn.relu'}), '(inputs=pool2_flat, units=1024, activation=tf.nn.relu)\n', (2823, 2877), True, 'import tensorflow as tf\n'), ((2964, 3056), 'tensorflow.layers.dropout', 'tf.layers.dropout', ([], {'inputs': 'dense', 'rate': '(0.4)', 'training': '(mode == tf.estimator.ModeKeys.TRAIN)'}), '(inputs=dense, rate=0.4, training=mode == tf.estimator.\n    ModeKeys.TRAIN)\n', (2981, 3056), True, 'import tensorflow as tf\n'), ((3183, 3224), 'tensorflow.layers.dense', 'tf.layers.dense', ([], {'inputs': 'dropout', 'units': '(10)'}), '(inputs=dropout, units=10)\n', (3198, 3224), True, 'import tensorflow as tf\n'), ((3727, 3795), 'tensorflow.losses.sparse_softmax_cross_entropy', 'tf.losses.sparse_softmax_cross_entropy', ([], {'labels': 'labels', 'logits': 'logits'}), '(labels=labels, logits=logits)\n', (3765, 3795), True, 'import tensorflow as tf\n'), ((4348, 4434), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'loss': 'loss', 'eval_metric_ops': 'eval_metric_ops'}), '(mode=mode, loss=loss, eval_metric_ops=\n    eval_metric_ops)\n', (4374, 4434), True, 'import tensorflow as tf\n'), ((4856, 4881), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (4879, 4881), False, 'import argparse\n'), ((5596, 5652), 'tensorflow.estimator.export.ServingInputReceiver', 'tf.estimator.export.ServingInputReceiver', (['inputs', 'inputs'], {}), '(inputs, inputs)\n', (5636, 5652), True, 'import tensorflow as tf\n'), ((5885, 5956), 'tensorflow.estimator.Estimator', 'tf.estimator.Estimator', ([], {'model_fn': 'cnn_model_fn', 'model_dir': 'args.model_dir'}), '(model_fn=cnn_model_fn, model_dir=args.model_dir)\n', (5907, 5956), True, 'import tensorflow as tf\n'), ((6143, 6210), 'tensorflow.train.LoggingTensorHook', 'tf.train.LoggingTensorHook', ([], {'tensors': 'tensors_to_log', 'every_n_iter': '(50)'}), '(tensors=tensors_to_log, every_n_iter=50)\n', (6169, 6210), True, 'import tensorflow as tf\n'), ((6255, 6377), 'tensorflow.estimator.inputs.numpy_input_fn', 'tf.estimator.inputs.numpy_input_fn', ([], {'x': "{'x': train_data}", 'y': 'train_labels', 'batch_size': '(100)', 'num_epochs': 'None', 'shuffle': '(True)'}), "(x={'x': train_data}, y=train_labels,\n    batch_size=100, num_epochs=None, shuffle=True)\n", (6289, 6377), True, 'import tensorflow as tf\n'), ((6484, 6586), 'tensorflow.estimator.inputs.numpy_input_fn', 'tf.estimator.inputs.numpy_input_fn', ([], {'x': "{'x': eval_data}", 'y': 'eval_labels', 'num_epochs': '(1)', 'shuffle': '(False)'}), "(x={'x': eval_data}, y=eval_labels,\n    num_epochs=1, shuffle=False)\n", (6518, 6586), True, 'import tensorflow as tf\n'), ((6639, 6694), 'tensorflow.estimator.TrainSpec', 'tf.estimator.TrainSpec', (['train_input_fn'], {'max_steps': '(20000)'}), '(train_input_fn, max_steps=20000)\n', (6661, 6694), True, 'import tensorflow as tf\n'), ((6711, 6747), 'tensorflow.estimator.EvalSpec', 'tf.estimator.EvalSpec', (['eval_input_fn'], {}), '(eval_input_fn)\n', (6732, 6747), True, 'import tensorflow as tf\n'), ((6752, 6824), 'tensorflow.estimator.train_and_evaluate', 'tf.estimator.train_and_evaluate', (['mnist_classifier', 'train_spec', 'eval_spec'], {}), '(mnist_classifier, train_spec, eval_spec)\n', (6783, 6824), True, 'import tensorflow as tf\n'), ((3324, 3355), 'tensorflow.argmax', 'tf.argmax', ([], {'input': 'logits', 'axis': '(1)'}), '(input=logits, axis=1)\n', (3333, 3355), True, 'import tensorflow as tf\n'), ((3487, 3531), 'tensorflow.nn.softmax', 'tf.nn.softmax', (['logits'], {'name': '"""softmax_tensor"""'}), "(logits, name='softmax_tensor')\n", (3500, 3531), True, 'import tensorflow as tf\n'), ((3599, 3661), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'predictions': 'predictions'}), '(mode=mode, predictions=predictions)\n', (3625, 3661), True, 'import tensorflow as tf\n'), ((3910, 3964), 'tensorflow.train.GradientDescentOptimizer', 'tf.train.GradientDescentOptimizer', ([], {'learning_rate': '(0.001)'}), '(learning_rate=0.001)\n', (3943, 3964), True, 'import tensorflow as tf\n'), ((4094, 4161), 'tensorflow.estimator.EstimatorSpec', 'tf.estimator.EstimatorSpec', ([], {'mode': 'mode', 'loss': 'loss', 'train_op': 'train_op'}), '(mode=mode, loss=loss, train_op=train_op)\n', (4120, 4161), True, 'import tensorflow as tf\n'), ((4252, 4322), 'tensorflow.metrics.accuracy', 'tf.metrics.accuracy', ([], {'labels': 'labels', 'predictions': "predictions['classes']"}), "(labels=labels, predictions=predictions['classes'])\n", (4271, 4322), True, 'import tensorflow as tf\n'), ((4498, 4538), 'os.path.join', 'os.path.join', (['base_dir', '"""train_data.npy"""'], {}), "(base_dir, 'train_data.npy')\n", (4510, 4538), False, 'import os\n'), ((4562, 4604), 'os.path.join', 'os.path.join', (['base_dir', '"""train_labels.npy"""'], {}), "(base_dir, 'train_labels.npy')\n", (4574, 4604), False, 'import os\n'), ((4691, 4730), 'os.path.join', 'os.path.join', (['base_dir', '"""eval_data.npy"""'], {}), "(base_dir, 'eval_data.npy')\n", (4703, 4730), False, 'import os\n'), ((4753, 4794), 'os.path.join', 'os.path.join', (['base_dir', '"""eval_labels.npy"""'], {}), "(base_dir, 'eval_labels.npy')\n", (4765, 4794), False, 'import os\n'), ((5544, 5583), 'tensorflow.placeholder', 'tf.placeholder', (['tf.float32', '[None, 784]'], {}), '(tf.float32, [None, 784])\n', (5558, 5583), True, 'import tensorflow as tf\n'), ((5148, 5178), 'os.environ.get', 'os.environ.get', (['"""SM_MODEL_DIR"""'], {}), "('SM_MODEL_DIR')\n", (5162, 5178), False, 'import os\n'), ((5233, 5270), 'os.environ.get', 'os.environ.get', (['"""SM_CHANNEL_TRAINING"""'], {}), "('SM_CHANNEL_TRAINING')\n", (5247, 5270), False, 'import os\n'), ((5426, 5459), 'os.environ.get', 'os.environ.get', (['"""SM_CURRENT_HOST"""'], {}), "('SM_CURRENT_HOST')\n", (5440, 5459), False, 'import os\n'), ((4051, 4077), 'tensorflow.train.get_global_step', 'tf.train.get_global_step', ([], {}), '()\n', (4075, 4077), True, 'import tensorflow as tf\n'), ((5337, 5363), 'os.environ.get', 'os.environ.get', (['"""SM_HOSTS"""'], {}), "('SM_HOSTS')\n", (5351, 5363), False, 'import os\n')]