Remove TFOptimizer wrapper

felker · felker · commit 64a9c77b6a1d · 2020-06-18T17:33:05.000-05:00
diff --git a/plasma/models/mpi_runner.py b/plasma/models/mpi_runner.py
@@ -255,21 +255,24 @@ def compile(self, optimizer, clipnorm, loss='mse'):
         # TODO(KGF): check the following import taken from runner.py
         # Was not in this file, originally.
         from tensorflow.keras.optimizers import (
-            SGD, Adam, RMSprop, Nadam, TFOptimizer
+            SGD, Adam, RMSprop, Nadam
             )
         if optimizer == 'sgd':
             optimizer_class = SGD(lr=self.DUMMY_LR, clipnorm=clipnorm)
         elif optimizer == 'momentum_sgd':
             optimizer_class = SGD(lr=self.DUMMY_LR, clipnorm=clipnorm,
                                   decay=1e-6, momentum=0.9)
         elif optimizer == 'tf_momentum_sgd':
-            optimizer_class = TFOptimizer(tf.train.MomentumOptimizer(
-                learning_rate=self.DUMMY_LR, momentum=0.9))
+            # TODO(KGF): removed TFOptimizer wrapper from here and below
+            # may not work anymore? See
+            # https://github.com/tensorflow/tensorflow/issues/22780
+            optimizer_class = tf.train.MomentumOptimizer(
+                learning_rate=self.DUMMY_LR, momentum=0.9)
         elif optimizer == 'adam':
             optimizer_class = Adam(lr=self.DUMMY_LR, clipnorm=clipnorm)
         elif optimizer == 'tf_adam':
-            optimizer_class = TFOptimizer(tf.train.AdamOptimizer(
-                learning_rate=self.DUMMY_LR))
+            optimizer_class = tf.train.AdamOptimizer(
+                learning_rate=self.DUMMY_LR)
         elif optimizer == 'rmsprop':
             optimizer_class = RMSprop(lr=self.DUMMY_LR, clipnorm=clipnorm)
         elif optimizer == 'nadam':
diff --git a/plasma/utils/state_reset.py b/plasma/utils/state_reset.py
@@ -7,7 +7,7 @@ def get_states(model):
         if hasattr(layer, "states"):
             layer_states = []
             for state in layer.states:
-                import keras.backend as K
+                import tensorflow.keras.backend as K
                 layer_states.append(K.get_value(state))
             all_states.append(layer_states)
     return all_states
