Reduce usage of tf.contrib (tensorflow#1345)

* tf.contrib.rnn -> tf.nn.rnn_cell

* tf.contrib.framework -> tf.train

tf.contrib.framework.load_checkpoint -> tf.train.load_checkpoint
tf.contrib.framework.list_variables -> tf.train.list_variables
tf.contrib.framework.init_from_checkpoint -> tf.train.init_from_checkpoint

* tf.contrib.learn.ModeKeys -> tf.estimator.ModeKeys

tf.contrib.learn.ModeKeys.TRAIN -> tf.estimator.ModeKeys.TRAIN
tf.contrib.learn.ModeKeys.EVAL -> tf.estimator.ModeKeys.EVAL
tf.contrib.learn.ModeKeys.INFER -> tf.estimator.ModeKeys.PREDICT

* tf.contrib.eager.in_eager_mode -> tf.executing_eagerly

* tf.contrib.data.batch_and_drop_remainder -> .batch(..., drop_remainder=True)

* tf.contrib.data -> tf.data.experimental

* tf.contrib.util.make_tensor_proto -> tf.make_tensor_proto

* Simplify tf.contrib.eager imports

* tf.initialize_all_variables -> tf.global_variables_initializer

Author: lgeiger

tensor2tensor/bin/t2t_attack.py

@@ -200,7 +200,7 @@ def main(argv):
     sur_ch_model.get_probs(inputs)
 
     checkpoint_path = os.path.expanduser(FLAGS.surrogate_output_dir)
-    tf.contrib.framework.init_from_checkpoint(
+    tf.train.init_from_checkpoint(
         tf.train.latest_checkpoint(checkpoint_path), {"/": "surrogate/"})
     sess.run(tf.global_variables_initializer())
 

tensor2tensor/bin/t2t_avg_all.py

@@ -60,7 +60,7 @@ def main(_):
   for model in bleu_hook.stepfiles_iterator(model_dir, FLAGS.wait_minutes,
                                             FLAGS.min_steps):
     if models_processed == 0:
-      var_list = tf.contrib.framework.list_variables(model.filename)
+      var_list = tf.train.list_variables(model.filename)
       avg_values = {}
       for (name, shape) in var_list:
         if not (name.startswith("global_step") or
@@ -69,7 +69,7 @@ def main(_):
     models_processed += 1
 
     tf.logging.info("Loading [%d]: %s" % (models_processed, model.filename))
-    reader = tf.contrib.framework.load_checkpoint(model.filename)
+    reader = tf.train.load_checkpoint(model.filename)
     for name in avg_values:
       avg_values[name] += reader.get_tensor(name) / FLAGS.n
     queue.append(model)
@@ -106,7 +106,7 @@ def main(_):
     tf.reset_default_graph()
     first_model = queue.popleft()
 
-    reader = tf.contrib.framework.load_checkpoint(first_model.filename)
+    reader = tf.train.load_checkpoint(first_model.filename)
     for name in avg_values:
       avg_values[name] -= reader.get_tensor(name) / FLAGS.n
 

tensor2tensor/data_generators/image_utils.py

@@ -262,7 +262,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 def encode_images_as_png(images):
   """Yield images encoded as pngs."""
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     for image in images:
       yield tf.image.encode_png(image).numpy()
   else:

tensor2tensor/data_generators/problem.py

@@ -423,7 +423,7 @@ def _preprocess(example):
 
     if interleave:
       dataset = dataset.apply(
-          tf.contrib.data.parallel_interleave(
+          tf.data.experimental.parallel_interleave(
               _preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = dataset.flat_map(_preprocess)
@@ -674,7 +674,7 @@ def _load_records_and_preprocess(filenames):
     # Create data-set from files by parsing, pre-processing and interleaving.
     if shuffle_files:
       dataset = dataset.apply(
-          tf.contrib.data.parallel_interleave(
+          tf.data.experimental.parallel_interleave(
               _load_records_and_preprocess, sloppy=True, cycle_length=8))
     else:
       dataset = _load_records_and_preprocess(dataset)
@@ -963,7 +963,7 @@ def define_shapes(example):
           batching_scheme["batch_sizes"] = [hparams.batch_size]
           batching_scheme["boundaries"] = []
         dataset = dataset.apply(
-            tf.contrib.data.bucket_by_sequence_length(
+            tf.data.experimental.bucket_by_sequence_length(
                 data_reader.example_length, batching_scheme["boundaries"],
                 batching_scheme["batch_sizes"]))
 
@@ -1040,7 +1040,7 @@ def serving_input_fn(self, hparams):
         tf.shape(serialized_example, out_type=tf.int64)[0],
         dataset.output_shapes)
     dataset = dataset.map(standardize_shapes)
-    features = tf.contrib.data.get_single_element(dataset)
+    features = tf.data.experimental.get_single_element(dataset)
 
     if self.has_inputs:
       features.pop("targets", None)

tensor2tensor/layers/common_image_attention.py

@@ -501,7 +501,7 @@ def postprocess_image(x, rows, cols, hparams):
                               use_bias=True,
                               activation=None,
                               name="output_conv")
-  if (hparams.mode == tf.contrib.learn.ModeKeys.INFER and
+  if (hparams.mode == tf.estimator.ModeKeys.PREDICT and
       hparams.block_raster_scan):
     y = targets
     yshape = common_layers.shape_list(y)
@@ -547,7 +547,7 @@ def prepare_decoder(targets, hparams):
 
   # during training, images are [batch, IMG_LEN, IMG_LEN, 3].
   # At inference, they are [batch, curr_infer_length, 1, 1]
-  if hparams.mode == tf.contrib.learn.ModeKeys.INFER:
+  if hparams.mode == tf.estimator.ModeKeys.PREDICT:
     curr_infer_length = targets_shape[1]
     if hparams.block_raster_scan:
       assert hparams.img_len*channels % hparams.query_shape[1] == 0

tensor2tensor/layers/common_image_attention_test.py

@@ -61,7 +61,7 @@ def testPostProcessImageInferMode(self, likelihood, num_mixtures, depth):
         block_raster_scan=True,
         hidden_size=2,
         likelihood=likelihood,
-        mode=tf.contrib.learn.ModeKeys.INFER,
+        mode=tf.estimator.ModeKeys.PREDICT,
         num_mixtures=num_mixtures,
         query_shape=[block_length, block_width],
     )

tensor2tensor/layers/common_layers.py

@@ -349,7 +349,7 @@ def embedding(x,
     # On the backwards pass, we want to convert the gradient from
     # an indexed-slices to a regular tensor before sending it back to the
     # parameter server. This avoids excess computation on the parameter server.
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       embedding_var = convert_gradient_to_tensor(embedding_var)
     x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
     emb_x = gather(embedding_var, x, dtype)
@@ -2868,7 +2868,7 @@ def ones_matrix_band_part(rows, cols, num_lower, num_upper, out_shape=None):
 def reshape_like_all_dims(a, b):
   """Reshapes a to match the shape of b."""
   ret = tf.reshape(a, tf.shape(b))
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape(b.get_shape())
   return ret
 
@@ -3193,7 +3193,7 @@ def should_generate_summaries():
 def reshape_like(a, b):
   """Reshapes a to match the shape of b in all but the last dimension."""
   ret = tf.reshape(a, tf.concat([tf.shape(b)[:-1], tf.shape(a)[-1:]], 0))
-  if not tf.contrib.eager.in_eager_mode():
+  if not tf.executing_eagerly():
     ret.set_shape(b.get_shape().as_list()[:-1] + a.get_shape().as_list()[-1:])
   return ret
 
@@ -3205,7 +3205,7 @@ def summarize_video(video, prefix, max_outputs=1):
     raise ValueError("Assuming videos given as tensors in the format "
                      "[batch, time, height, width, channels] but got one "
                      "of shape: %s" % str(video_shape))
-  if tf.contrib.eager.in_eager_mode():
+  if tf.executing_eagerly():
     return
   if video.get_shape().as_list()[1] is None:
     tf.summary.image(

tensor2tensor/layers/common_video.py

@@ -82,7 +82,7 @@ def lstm_cell(inputs,
               name=None):
   """Full LSTM cell."""
   input_shape = common_layers.shape_list(inputs)
-  cell = tf.contrib.rnn.LSTMCell(num_units,
+  cell = tf.nn.rnn_cell.LSTMCell(num_units,
                                  use_peepholes=use_peepholes,
                                  cell_clip=cell_clip,
                                  initializer=initializer,

tensor2tensor/layers/discretization.py

@@ -473,7 +473,7 @@ def gumbel_softmax(x,
     d_dev = -tf.reduce_mean(d_variance)
     ret = s
 
-    if mode != tf.contrib.learn.ModeKeys.TRAIN:
+    if mode != tf.estimator.ModeKeys.TRAIN:
       ret = tf.reshape(maxvhot, common_layers.shape_list(s))  # Just hot @eval.
     return m, ret, d_dev * 5.0 + tf.reduce_mean(kl) * 0.002
 
@@ -822,7 +822,7 @@ def predict_bits_with_lstm(prediction_source, state_size, total_num_bits,
 
   with tf.variable_scope("predict_bits_with_lstm"):
     # Layers and cell state creation.
-    lstm_cell = tf.contrib.rnn.LSTMCell(state_size)
+    lstm_cell = tf.nn.rnn_cell.LSTMCell(state_size)
     discrete_predict = tf.layers.Dense(2**bits_at_once, name="discrete_predict")
     discrete_embed = tf.layers.Dense(state_size, name="discrete_embed")
     batch_size = common_layers.shape_list(prediction_source)[0]

tensor2tensor/layers/modalities.py

@@ -93,7 +93,7 @@ def _get_weights(self, hidden_dim=None):
     else:
       ret = tf.concat(shards, 0)
     # Convert ret to tensor.
-    if not tf.contrib.eager.in_eager_mode():
+    if not tf.executing_eagerly():
       ret = common_layers.convert_gradient_to_tensor(ret)
     return ret
 
@@ -226,15 +226,15 @@ class ImageModality(modality.Modality):
 
   def bottom(self, x):
     with tf.variable_scope(self.name):
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         tf.summary.image(
             "inputs", common_layers.tpu_safe_image_summary(x), max_outputs=2)
       return tf.to_float(x)
 
   def targets_bottom(self, x):
     inputs = x
     with tf.variable_scope(self.name):
-      if not tf.contrib.eager.in_eager_mode():
+      if not tf.executing_eagerly():
         tf.summary.image(
             "targets_bottom",
             common_layers.tpu_safe_image_summary(inputs),