From 582b71131d75a1804c306b602959518797c0a01b Mon Sep 17 00:00:00 2001
From: nud000j <nud000j@walmartlabs.com>
Date: Tue, 15 Oct 2019 16:13:05 -0400
Subject: [PATCH] Fix keras imports

---
 contrib/tf_serving/save_tfs_model.py    |  6 +++---
 contrib/tf_serving/tfs_sample_client.py |  2 +-
 src/handlers/data_generator.py          |  2 +-
 src/handlers/model_builder.py           | 12 ++++++------
 src/trainer/train.py                    |  4 ++--
 src/utils/keras_utils.py                |  4 ++--
 src/utils/losses.py                     |  2 +-
 src/utils/utils.py                      |  2 +-
 8 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/contrib/tf_serving/save_tfs_model.py b/contrib/tf_serving/save_tfs_model.py
index d177b5cc..82042bc2 100644
--- a/contrib/tf_serving/save_tfs_model.py
+++ b/contrib/tf_serving/save_tfs_model.py
@@ -1,7 +1,7 @@
-import keras.backend as K
+import tensorflow.keras.backend as K
 import argparse
-from keras.applications.mobilenet import DepthwiseConv2D, relu6
-from keras.utils.generic_utils import CustomObjectScope
+from tensorflow.keras.applications.mobilenet import DepthwiseConv2D, relu6
+from tensorflow.keras.utils.generic_utils import CustomObjectScope
 from tensorflow.python.saved_model import builder as saved_model_builder
 from tensorflow.python.saved_model import tag_constants
 from tensorflow.python.saved_model.signature_def_utils_impl import \
diff --git a/contrib/tf_serving/tfs_sample_client.py b/contrib/tf_serving/tfs_sample_client.py
index 3055bf78..5208db9d 100644
--- a/contrib/tf_serving/tfs_sample_client.py
+++ b/contrib/tf_serving/tfs_sample_client.py
@@ -1,6 +1,6 @@
 import json
 import argparse
-import keras
+import tensorflow.keras as keras
 import numpy as np
 import tensorflow as tf
 from src.utils import utils
diff --git a/src/handlers/data_generator.py b/src/handlers/data_generator.py
index 65e772fe..4736fbd8 100644
--- a/src/handlers/data_generator.py
+++ b/src/handlers/data_generator.py
@@ -1,7 +1,7 @@
 
 import os
 import numpy as np
-import keras
+import tensorflow.keras as keras
 from utils import utils
 
 
diff --git a/src/handlers/model_builder.py b/src/handlers/model_builder.py
index ccbcef78..ee15276e 100644
--- a/src/handlers/model_builder.py
+++ b/src/handlers/model_builder.py
@@ -1,8 +1,8 @@
 
 import importlib
-from keras.models import Model
-from keras.layers import Dropout, Dense
-from keras.optimizers import Adam
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Dropout, Dense
+from tensorflow.keras.optimizers import Adam
 from utils.losses import earth_movers_distance
 
 
@@ -21,11 +21,11 @@ def __init__(self, base_model_name, n_classes=10, learning_rate=0.001, dropout_r
     def _get_base_module(self):
         # import Keras base model module
         if self.base_model_name == 'InceptionV3':
-            self.base_module = importlib.import_module('keras.applications.inception_v3')
+            self.base_module = importlib.import_module('tensorflow.keras.applications.inception_v3')
         elif self.base_model_name == 'InceptionResNetV2':
-            self.base_module = importlib.import_module('keras.applications.inception_resnet_v2')
+            self.base_module = importlib.import_module('tensorflow.keras.applications.inception_resnet_v2')
         else:
-            self.base_module = importlib.import_module('keras.applications.'+self.base_model_name.lower())
+            self.base_module = importlib.import_module('tensorflow.keras.applications.'+self.base_model_name.lower())
 
     def build(self):
         # get base model class
diff --git a/src/trainer/train.py b/src/trainer/train.py
index 0ba8b7ba..ae72a6bc 100644
--- a/src/trainer/train.py
+++ b/src/trainer/train.py
@@ -1,8 +1,8 @@
 
 import os
 import argparse
-from keras import backend as K
-from keras.callbacks import ModelCheckpoint
+from tensorflow.keras import backend as K
+from tensorflow.keras.callbacks import ModelCheckpoint
 from sklearn.model_selection import train_test_split
 from handlers.data_generator import TrainDataGenerator, TestDataGenerator
 from handlers.model_builder import Nima
diff --git a/src/utils/keras_utils.py b/src/utils/keras_utils.py
index 6b986a34..15747e86 100644
--- a/src/utils/keras_utils.py
+++ b/src/utils/keras_utils.py
@@ -1,6 +1,6 @@
 
-import keras.backend as K
-from keras.callbacks import TensorBoard
+import tensorflow.keras.backend as K
+from tensorflow.keras.callbacks import TensorBoard
 
 
 class TensorBoardBatch(TensorBoard):
diff --git a/src/utils/losses.py b/src/utils/losses.py
index a9454f88..cf78ab26 100644
--- a/src/utils/losses.py
+++ b/src/utils/losses.py
@@ -1,5 +1,5 @@
 
-from keras import backend as K
+from tensorflow.keras import backend as K
 
 
 def earth_movers_distance(y_true, y_pred):
diff --git a/src/utils/utils.py b/src/utils/utils.py
index 789dc8cb..559f2c09 100644
--- a/src/utils/utils.py
+++ b/src/utils/utils.py
@@ -1,7 +1,7 @@
 
 import os
 import json
-import keras
+import tensorflow.keras as keras
 import numpy as np