Save Tensorflow model to Minio repository
[demo.git] / vnfs / DAaaS / applications / sample-horovod-app / keras_mnist_advanced_modified.py
index 03425ff..fa39cb6 100644 (file)
@@ -1,13 +1,19 @@
 from __future__ import print_function
 import keras
-from keras.datasets import mnist
-from keras.models import Sequential
-from keras.layers import Dense, Dropout, Flatten
-from keras.layers import Conv2D, MaxPooling2D
-from keras.preprocessing.image import ImageDataGenerator
-from keras import backend as K
+import os
+from tensorflow.keras.datasets import mnist
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Dropout, Flatten
+from tensorflow.keras.layers import Conv2D, MaxPooling2D
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras import backend as K
+from tensorflow_estimator.python.estimator.export import export as export_helpers
+from tensorflow.python.saved_model import builder as saved_model_builder
+from tensorflow.python.saved_model import tag_constants, signature_constants
+from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def
 import tensorflow as tf
-import horovod.keras as hvd
+import horovod.tensorflow.keras as hvd
+
 
 # Horovod: initialize Horovod.
 hvd.init()
@@ -53,8 +59,8 @@ print(x_train.shape[0], 'train samples')
 print(x_test.shape[0], 'test samples')
 
 # Convert class vectors to binary class matrices
-y_train = keras.utils.to_categorical(y_train, num_classes)
-y_test = keras.utils.to_categorical(y_test, num_classes)
+y_train = tf.keras.utils.to_categorical(y_train, num_classes)
+y_test = tf.keras.utils.to_categorical(y_test, num_classes)
 
 model = Sequential()
 model.add(Conv2D(32, kernel_size=(3, 3),
@@ -69,12 +75,12 @@ model.add(Dropout(0.5))
 model.add(Dense(num_classes, activation='softmax'))
 
 # Horovod: adjust learning rate based on number of GPUs.
-opt = keras.optimizers.Adadelta(lr=1.0 * hvd.size())
+opt = tf.keras.optimizers.Adadelta(lr=1.0 * hvd.size())
 
 # Horovod: add Horovod Distributed Optimizer.
 opt = hvd.DistributedOptimizer(opt)
 
-model.compile(loss=keras.losses.categorical_crossentropy,
+model.compile(loss=tf.keras.losses.categorical_crossentropy,
               optimizer=opt,
               metrics=['accuracy'])
 
@@ -96,12 +102,13 @@ callbacks = [
     hvd.callbacks.LearningRateWarmupCallback(warmup_epochs=5, verbose=1),
 
     # Reduce the learning rate if training plateaues.
-    keras.callbacks.ReduceLROnPlateau(patience=10, verbose=1),
+    tf.keras.callbacks.ReduceLROnPlateau(patience=10, verbose=1),
 ]
 
 # Horovod: save checkpoints only on worker 0 to prevent other workers from corrupting them.
 if hvd.rank() == 0:
-    callbacks.append(keras.callbacks.ModelCheckpoint('./checkpoint-{epoch}.h5'))
+    callbacks.append(tf.keras.callbacks.ModelCheckpoint(
+        './checkpoint-{epoch}.h5'))
 
 # Set up ImageDataGenerators to do data augmentation for the training images.
 train_gen = ImageDataGenerator(rotation_range=8, width_shift_range=0.08, shear_range=0.3,
@@ -118,10 +125,45 @@ model.fit_generator(train_gen.flow(x_train, y_train, batch_size=batch_size),
                     callbacks=callbacks,
                     epochs=epochs,
                     verbose=1,
-                    validation_data=test_gen.flow(x_test, y_test, batch_size=batch_size),
+                    validation_data=test_gen.flow(
+                        x_test, y_test, batch_size=batch_size),
                     validation_steps=3 * test_batches // hvd.size())
 
 # Evaluate the model on the full data set.
 score = model.evaluate(x_test, y_test, verbose=0)
 print('Test loss:', score[0])
 print('Test accuracy:', score[1])
+
+# Save Model to Minio
+if hvd.rank() == 0:
+    print('Model Summary')
+    model.summary()
+    print('Exporting trained model to Minio Model Repo')
+    base_path = os.environ['MODEL_BASE_PATH']
+
+    # Option 1(Preferred) - Using Keras api and Tensorflow v1.13 version
+    saved_model_path = tf.contrib.saved_model.save_keras_model(model, base_path)
+    print('Model Saved to {} Using new Keras API!!!'.format(saved_model_path))
+    # Option 2 - Tensorflow v1.13+ Builder saved_model api.
+    # builder = saved_model_builder.SavedModelBuilder(base_path)
+
+    # print(model.input)
+    # print(model.outputs)
+
+    # signature = predict_signature_def(inputs={"inputs": model.input},
+    #                                   outputs={t.name:t for t in model.outputs})
+    # print(signature)
+    # K.set_learning_phase(0)
+    # with K.get_session() as sess:
+    #     builder.add_meta_graph_and_variables(sess=sess,
+    #                                          tags=[tag_constants.SERVING],
+    #                                          signature_def_map={'predict': signature})
+    #     builder.save()
+    # print('Model Saved to S3 Using Builder!!!')
+
+    # Option 3 - Tensorflow v1.13 Will be deprecated in Tensorflow v2
+    # tf.saved_model.simple_save(
+    #     keras.backend.get_session(),
+    #     export_path,
+    #     inputs={'input_image': model.input},
+    #     outputs={t.name: t for t in model.outputs})