#!/usr/bin/env python3

import sys

import time

import tensorflow as tf

from tensorflow import keras

import numpy as np

from datetime import datetime

import json

import os

import boto3

import io

print('Arguments received:', str(sys.argv))

print('Env variables', os.environ)

s3 = boto3.client('s3')

cloudwatch_client = boto3.client('cloudwatch', region_name='us-east-1')

s3_bucket = sys.argv[1]

str_keras_model_path = '/app/mnist_model.h5'

str_tflite_model_path = '/app/mnist_model.tflite'

## Importing hyperparameters

result = s3.get_object(Bucket=s3_bucket, Key=sys.argv[2])

hyperparameters = json.loads(result["Body"].read().decode())

print(hyperparameters)

parameters = {}

parameters['input'] = hyperparameters

fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

train_images = train_images / 255.0

test_images = test_images / 255.0

model = keras.Sequential([

keras.layers.Flatten(input_shape=(28, 28)),

keras.layers.Dense(128, activation=tf.nn.relu),

keras.layers.Dense(10, activation=tf.nn.softmax)

])

model.compile(optimizer='adam',

loss='sparse_categorical_crossentropy',

metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=int(hyperparameters['num_of_epochs']))

test_loss, test_acc = model.evaluate(test_images, test_labels)

print('Test accuracy:', test_acc)

predictions = model.predict(test_images)

print(np.argmax(predictions[0]))

print(test_labels[0])

print("Accuracy = {}".format(test_acc))

cloudwatch_client.put_metric_data(

Namespace='production',

MetricData=[

{

'MetricName' : 'Accuracy',

'Dimensions' : [

{

'Name' : 'project',

'Value' : 'training'

}

],

'Timestamp': datetime.now(),

'Value': test_acc

}

]

)

# Save the Keras model

model.save(str_keras_model_path)

# Convert the Keras model to TFLite model

converter = tf.contrib.lite.TFLiteConverter.from_keras_model_file(

str_keras_model_path, input_shapes={'input_1' : [1,299,299,3]}

)

tflite_model = converter.convert()

open(str_tflite_model_path, "wb").write(tflite_model)

# Upload the TFLite model to S3

with open(str_tflite_model_path, 'rb') as data:

s3.upload_fileobj(data, s3_bucket, sys.argv[4])

parameters['output'] = {

'accuracy' : test_acc,

'model_path' : sys.argv[4]

}

print(parameters)

s3.upload_fileobj(io.BytesIO(json.dumps(parameters).encode('utf-8')), s3_bucket, sys.argv[3])

sys.exit(0)