Remove tf.compat.v1 imports in model/

Some minor API changes, but not many. We do have to use some carve outs
where we manually write out tf.compat.v1 before some API usages, but I
expect these to go away as we hack more on the code base and fix more
things up.

Reviewers: orodley, ondrasej, virajbshah

Reviewed By: ondrasej

Pull Request: #332

Author: boomanaiden154

gematria/model/python/inference.py

@@ -21,7 +21,6 @@
 from gematria.model.python import model_base
 from gematria.model.python import training
 from gematria.proto import throughput_pb2
-import tensorflow.compat.v1 as tf
 
 
 def _get_num_instructions_in_block_with_throughput_proto(

gematria/model/python/loss_utils.py

@@ -16,7 +16,7 @@
 from collections.abc import Sequence
 
 from gematria.model.python import options
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 import tensorflow_probability as tfp
 
 

gematria/model/python/main_function.py

@@ -47,8 +47,7 @@ def main(_):
 from gematria.proto import throughput_pb2
 from gematria.utils.python import timer
 import numpy as np
-import tensorflow.compat.v1 as tf
-import tensorflow as tf2
+import tensorflow as tf
 
 _ACTION = flags.DEFINE_enum_class(
     'gematria_action',
@@ -685,7 +684,7 @@ def _restore_model_from_checkpoint(
     checkpoint_file: str, model: model_base.ModelBase
 ) -> None:
   """Restores a model from a checkpoint."""
-  checkpoint = tf2.train.Checkpoint(model)
+  checkpoint = tf.train.Checkpoint(model)
   checkpoint.restore(checkpoint_file)
 
 
@@ -720,10 +719,10 @@ def run_gematria_model_from_command_line_flags(
   """
   np.set_printoptions(edgeitems=_GEMATRIA_NUMPY_PRINT_EDGEITEMS.value)
   if _GEMATRIA_RANDOM_SEED.value >= 0:
-    tf.random.set_random_seed(_GEMATRIA_RANDOM_SEED.value)
+    tf.random.set_seed(_GEMATRIA_RANDOM_SEED.value)
     random.seed(_GEMATRIA_RANDOM_SEED.value)
   is_chief = _GEMATRIA_TRAINING_TASK.value == 0
-  dev = tf.train.replica_device_setter(
+  dev = tf.compat.v1.train.replica_device_setter(
       ps_tasks=_GEMATRIA_TRAINING_PS_TASKS.value
   )
   with tf.device(dev):
@@ -812,8 +811,8 @@ def run_gematria_model_from_command_line_flags(
           == io_options.ThroughputSelection.RANDOM
       )
 
-      checkpoint = tf2.train.Checkpoint(model)
-      checkpoint_manager = tf2.train.CheckpointManager(
+      checkpoint = tf.train.Checkpoint(model)
+      checkpoint_manager = tf.train.CheckpointManager(
           checkpoint,
           _CHECKPOINT_DIR.value,
           _GEMATRIA_CHECKPOINT_MAX_TO_KEEP.value,
@@ -822,11 +821,11 @@ def run_gematria_model_from_command_line_flags(
       def checkpoint_model():
         checkpoint_manager.save()
 
-      train_summary_writer = tf2.summary.create_file_writer(
+      train_summary_writer = tf.summary.create_file_writer(
           _GEMATRIA_SUMMARY_DIR.value
       )
 
-      with train_summary_writer.as_default(), tf2.summary.record_if(
+      with train_summary_writer.as_default(), tf.summary.record_if(
           lambda: tf.math.equal(
               model.global_step % _GEMATRIA_SAVE_SUMMARIES_EPOCHS, 0
           )

gematria/model/python/model_base.py

@@ -43,7 +43,7 @@
 from gematria.utils.python import timer
 import numpy as np
 import scipy.stats
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 import tf_slim.evaluation
 
 # The type used for TensorFlow feed_dict objects. The type we use here is
@@ -75,7 +75,7 @@ class AddBasicBlockError(Exception):
   """The exception raised when adding a block to batch fails."""
 
 
-class SaveBestCheckpoint(tf.train.SessionRunHook):
+class SaveBestCheckpoint(tf.compat.v1.train.SessionRunHook):
   """A run hook that saves top N models based on error values."""
 
   def __init__(
@@ -100,8 +100,8 @@ def __init__(
     self._saver = tf.train.Saver(max_to_keep=max_to_keep, name='relative_mae')
     self._last_eval = math.inf
 
-  def before_run(self, run_context: ...) -> tf.train.SessionRunArgs:
-    return tf.train.SessionRunArgs(
+  def before_run(self, run_context: ...) -> tf.compat.v1.train.SessionRunArgs:
+    return tf.compat.v1.train.SessionRunArgs(
         {'loss': self._error_tensor, 'global_step': self._global_step}
     )
 
@@ -328,7 +328,7 @@ def __init__(
     )
     self._trained_variable_groups = tuple(trained_variable_groups or ())
 
-    self._global_step = tf.train.get_or_create_global_step()
+    self._global_step = tf.compat.v1.train.get_or_create_global_step()
 
     self._decayed_learning_rate = None
     self._loss: Optional[loss_utils.LossComputation] = None
@@ -604,61 +604,67 @@ def _create_optimizer(self) -> None:
           'must be great than zero.'
       )
     if self._learning_rate_schedule == options.LearningRateScheduleType.COSINE:
-      self._decayed_learning_rate = tf.train.cosine_decay(**decay_args)
+      self._decayed_learning_rate = tf.compat.v1.train.cosine_decay(
+          **decay_args
+      )
     elif (
         self._learning_rate_schedule
         == options.LearningRateScheduleType.EXPONENTIAL
     ):
-      self._decayed_learning_rate = tf.train.exponential_decay(
+      self._decayed_learning_rate = tf.compat.v1.train.exponential_decay(
           **decay_args, **decay_rate_arg
       )
     elif (
         self._learning_rate_schedule
         == options.LearningRateScheduleType.INVERSE_TIME
     ):
-      self._decayed_learning_rate = tf.train.inverse_time_decay(
+      self._decayed_learning_rate = tf.compat.v1.train.inverse_time_decay(
           **decay_args, **decay_rate_arg
       )
     elif (
         self._learning_rate_schedule
         == options.LearningRateScheduleType.LINEAR_COSINE
     ):
-      self._decayed_learning_rate = tf.train.linear_cosine_decay(**decay_args)
+      self._decayed_learning_rate = tf.compat.v1.train.linear_cosine_decay(
+          **decay_args
+      )
     elif (
         self._learning_rate_schedule
         == options.LearningRateScheduleType.NATURAL_EXP
     ):
-      self._decayed_learning_rate = tf.train.natural_exp_decay(
+      self._decayed_learning_rate = tf.compat.v1.train.natural_exp_decay(
           **decay_args, **decay_rate_arg
       )
     elif (
         self._learning_rate_schedule
         == options.LearningRateScheduleType.NOISY_LINEAR_COSINE
     ):
-      self._decayed_learning_rate = tf.train.noisy_linear_cosine_decay(
-          **decay_args
+      self._decayed_learning_rate = (
+          tf.compat.v1.train.noisy_linear_cosine_decay(**decay_args)
       )
     elif (
         self._learning_rate_schedule
         == options.LearningRateScheduleType.POLYNOMIAL
     ):
-      self._decayed_learning_rate = tf.train.polynomial_decay(**decay_args)
+      self._decayed_learning_rate = tf.compat.v1.train.polynomial_decay(
+          **decay_args
+      )
     else:
       assert (
           self._learning_rate_schedule == options.LearningRateScheduleType.NONE
       )
       self._decayed_learning_rate = self._learning_rate
 
     if self._optimizer_type == options.OptimizerType.ADAM:
-      self._optimizer = tf.train.AdamOptimizer(
+      self._optimizer = tf.compat.v1.train.AdamOptimizer(
           learning_rate=self._decayed_learning_rate
       )
     elif self._optimizer_type == options.OptimizerType.SGD:
-      self._optimizer = tf.train.GradientDescentOptimizer(
+      self._optimizer = tf.compat.v1.train.GradientDescentOptimizer(
           learning_rate=self._decayed_learning_rate
       )
     elif self._optimizer_type == options.OptimizerType.RMSPROP:
-      self._optimizer = tf.train.RMSPropOptimizer(
+      self._optimizer = tf.compat.v1.train.RMSPropOptimizer(
           learning_rate=self._decayed_learning_rate
       )
     else:
@@ -681,7 +687,7 @@ def _create_optimizer(self) -> None:
 
   def get_monitored_training_session_hooks(
       self,
-  ) -> Sequence[tf.train.SessionRunHook]:
+  ) -> Sequence[tf.compat.v1.train.SessionRunHook]:
     """Returns hooks for a MonitoredTrainingSession required by the model."""
     hooks = []
     if isinstance(self._optimizer, tf.train.SyncReplicasOptimizer):
@@ -1063,7 +1069,9 @@ def run_continuous_evaluation(
       tf_master: str = '',
       eval_interval_seconds: int = 60,
       max_num_evaluations: Optional[int] = None,
-      session_hooks: Optional[Sequence[tf.train.SessionRunHook]] = None,
+      session_hooks: Optional[
+          Sequence[tf.compat.v1.train.SessionRunHook]
+      ] = None,
       max_blocks_in_batch: Optional[int] = None,
       max_instructions_in_batch: Optional[int] = None,
   ) -> None:

gematria/model/python/model_blocks.py

@@ -17,7 +17,7 @@
 from typing import Optional
 
 import sonnet as snt
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 import tf_keras
 
 

gematria/model/python/token_model.py

@@ -25,7 +25,7 @@
 from gematria.model.python import model_base
 from gematria.model.python import oov_token_behavior
 import numpy as np
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 _OutOfVocabularyTokenBehavior = oov_token_behavior.OutOfVocabularyTokenBehavior
 

gematria/model/python/token_model_flags.py

@@ -22,7 +22,7 @@
 from absl import flags
 from gematria.model.python import oov_token_behavior
 from gematria.utils.python import flag_utils
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 _TOKEN_FILE = flags.DEFINE_string(
     'gematria_tokens_file',

gematria/model/python/token_model_flags_test.py

@@ -15,7 +15,7 @@
 from absl import flags
 from absl.testing import flagsaver
 from gematria.model.python import token_model_flags
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 FLAGS = flags.FLAGS
 

gematria/model/python/training.py

@@ -22,7 +22,7 @@
 from gematria.basic_block.python import basic_block
 from gematria.basic_block.python import throughput
 import numpy as np
-import tensorflow.compat.v1 as tf
+import tensorflow as tf
 
 
 @dataclasses.dataclass(frozen=True)