Minimal tensorflow 2 port

- Run `tf_upgrade_v2`
- Manually map some older tf.contrib objects not covered by `tf_upgrade_v2` (`MutableHashTable`, `layers.batch_norm`, `layers.dropout`.)
- Port the dropped `rnn_cell_impl._linear` function.
- Remove .value accessors for `get_shape()` values.

Author: mjpieters

aocr/__main__.py

@@ -17,7 +17,7 @@
 from .util.data_gen import DataGen
 from .util.export import Exporter
 
-tf.logging.set_verbosity(tf.logging.ERROR)
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
 
 
 def process_args(args, defaults):
@@ -212,7 +212,7 @@ def main(args=None):
     console.setFormatter(formatter)
     logging.getLogger('').addHandler(console)
 
-    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
+    with tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(allow_soft_placement=True)) as sess:
 
         if parameters.phase == 'dataset':
             dataset.generate(

aocr/model/cnn.py

@@ -15,10 +15,10 @@ def var_random(name, shape, regularizable=False):
     :param regularizable:
     :return:
     '''
-    v = tf.get_variable(name, shape=shape, initializer=tf.contrib.layers.xavier_initializer())
+    v = tf.compat.v1.get_variable(name, shape=shape, initializer=tf.compat.v1.keras.initializers.VarianceScaling(scale=1.0, mode="fan_avg", distribution="uniform"))
     if regularizable:
-        with tf.name_scope(name + '/Regularizer/'):
-            tf.add_to_collection(tf.GraphKeys.REGULARIZATION_LOSSES, tf.nn.l2_loss(v))
+        with tf.compat.v1.name_scope(name + '/Regularizer/'):
+            tf.compat.v1.add_to_collection(tf.compat.v1.GraphKeys.REGULARIZATION_LOSSES, tf.nn.l2_loss(v))
     return v
 
 
@@ -29,8 +29,8 @@ def max_2x2pool(incoming, name):
     :param name:
     :return:
     '''
-    with tf.variable_scope(name):
-        return tf.nn.max_pool(incoming, ksize=(1, 2, 2, 1), strides=(1, 2, 2, 1), padding='SAME')
+    with tf.compat.v1.variable_scope(name):
+        return tf.nn.max_pool2d(input=incoming, ksize=(1, 2, 2, 1), strides=(1, 2, 2, 1), padding='SAME')
 
 
 def max_2x1pool(incoming, name):
@@ -40,8 +40,8 @@ def max_2x1pool(incoming, name):
     :param name:
     :return:
     '''
-    with tf.variable_scope(name):
-        return tf.nn.max_pool(incoming, ksize=(1, 2, 1, 1), strides=(1, 2, 1, 1), padding='SAME')
+    with tf.compat.v1.variable_scope(name):
+        return tf.nn.max_pool2d(input=incoming, ksize=(1, 2, 1, 1), strides=(1, 2, 1, 1), padding='SAME')
 
 
 def ConvRelu(incoming, num_filters, filter_size, name):
@@ -54,14 +54,14 @@ def ConvRelu(incoming, num_filters, filter_size, name):
     :return:
     '''
     num_filters_from = incoming.get_shape().as_list()[3]
-    with tf.variable_scope(name):
+    with tf.compat.v1.variable_scope(name):
         conv_W = var_random(
             'W',
             tuple(filter_size) + (num_filters_from, num_filters),
             regularizable=True
         )
 
-        after_conv = tf.nn.conv2d(incoming, conv_W, strides=(1, 1, 1, 1), padding='SAME')
+        after_conv = tf.nn.conv2d(incoming, filters=conv_W, strides=(1, 1, 1, 1), padding='SAME')
 
         return tf.nn.relu(after_conv)
 
@@ -73,7 +73,7 @@ def batch_norm(incoming, is_training):
     :param is_training:
     :return:
     '''
-    return tf.contrib.layers.batch_norm(incoming, is_training=is_training, scale=True, decay=0.99)
+    return tf.compat.v1.layers.batch_normalization(incoming, training=is_training, scale=True, momentum=0.99)
 
 
 def ConvReluBN(incoming, num_filters, filter_size, name, is_training):
@@ -87,22 +87,22 @@ def ConvReluBN(incoming, num_filters, filter_size, name, is_training):
     :return:
     '''
     num_filters_from = incoming.get_shape().as_list()[3]
-    with tf.variable_scope(name):
+    with tf.compat.v1.variable_scope(name):
         conv_W = var_random(
             'W',
             tuple(filter_size) + (num_filters_from, num_filters),
             regularizable=True
         )
 
-        after_conv = tf.nn.conv2d(incoming, conv_W, strides=(1, 1, 1, 1), padding='SAME')
+        after_conv = tf.nn.conv2d(incoming, filters=conv_W, strides=(1, 1, 1, 1), padding='SAME')
 
         after_bn = batch_norm(after_conv, is_training)
 
         return tf.nn.relu(after_bn)
 
 
 def dropout(incoming, is_training, keep_prob=0.5):
-    return tf.contrib.layers.dropout(incoming, keep_prob=keep_prob, is_training=is_training)
+    return tf.compat.v1.layers.dropout(incoming, rate=1 - keep_prob, training=is_training)
 
 
 def tf_create_attention_map(incoming):

aocr/model/model.py

@@ -113,7 +113,7 @@ def __init__(self,
             self.height = tf.constant(DataGen.IMAGE_HEIGHT, dtype=tf.int32)
             self.height_float = tf.constant(DataGen.IMAGE_HEIGHT, dtype=tf.float64)
 
-            self.img_pl = tf.placeholder(tf.string, name='input_image_as_bytes')
+            self.img_pl = tf.compat.v1.placeholder(tf.string, name='input_image_as_bytes')
             self.img_data = tf.cond(
                 tf.less(tf.rank(self.img_pl), 1),
                 lambda: tf.expand_dims(self.img_pl, 0),
@@ -156,7 +156,7 @@ def __init__(self,
                 forward_only=self.forward_only,
                 use_gru=use_gru)
 
-            table = tf.contrib.lookup.MutableHashTable(
+            table = tf.lookup.experimental.MutableHashTable(
                 key_dtype=tf.int64,
                 value_dtype=tf.string,
                 default_value="",
@@ -226,12 +226,12 @@ def __init__(self,
                 self.updates = []
                 self.summaries_by_bucket = []
 
-                params = tf.trainable_variables()
-                opt = tf.train.AdadeltaOptimizer(learning_rate=initial_learning_rate)
+                params = tf.compat.v1.trainable_variables()
+                opt = tf.compat.v1.train.AdadeltaOptimizer(learning_rate=initial_learning_rate)
                 loss_op = self.attention_decoder_model.loss
 
                 if self.reg_val > 0:
-                    reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
+                    reg_losses = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.REGULARIZATION_LOSSES)
                     logging.info('Adding %s regularization losses', len(reg_losses))
                     logging.debug('REGULARIZATION_LOSSES: %s', reg_losses)
                     loss_op = self.reg_val * tf.reduce_sum(reg_losses) + loss_op
@@ -242,14 +242,14 @@ def __init__(self,
 
                 # Summaries for loss, variables, gradients, gradient norms and total gradient norm.
                 summaries = [
-                    tf.summary.scalar("loss", loss_op),
-                    tf.summary.scalar("total_gradient_norm", tf.global_norm(gradients))
+                    tf.compat.v1.summary.scalar("loss", loss_op),
+                    tf.compat.v1.summary.scalar("total_gradient_norm", tf.linalg.global_norm(gradients))
                 ]
-                all_summaries = tf.summary.merge(summaries)
+                all_summaries = tf.compat.v1.summary.merge(summaries)
                 self.summaries_by_bucket.append(all_summaries)
 
                 # update op - apply gradients
-                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
+                update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
                 with tf.control_dependencies(update_ops):
                     self.updates.append(
                         opt.apply_gradients(
@@ -258,7 +258,7 @@ def __init__(self,
                         )
                     )
 
-        self.saver_all = tf.train.Saver(tf.all_variables())
+        self.saver_all = tf.compat.v1.train.Saver(tf.compat.v1.all_variables())
         self.checkpoint_path = os.path.join(self.model_dir, "model.ckpt")
 
         ckpt = tf.train.get_checkpoint_state(model_dir)
@@ -268,7 +268,7 @@ def __init__(self,
             self.saver_all.restore(self.sess, ckpt.model_checkpoint_path)
         else:
             logging.info("Created model with fresh parameters.")
-            self.sess.run(tf.initialize_all_variables())
+            self.sess.run(tf.compat.v1.initialize_all_variables())
 
     def predict(self, image_file_data):
         input_feed = {}
@@ -370,7 +370,7 @@ def train(self, data_path, num_epoch):
         loss = 0.0
         current_step = 0
         skipped_counter = 0
-        writer = tf.summary.FileWriter(self.model_dir, self.sess.graph)
+        writer = tf.compat.v1.summary.FileWriter(self.model_dir, self.sess.graph)
 
         logging.info('Starting the training process.')
         for batch in s_gen.gen(self.batch_size):
@@ -498,18 +498,18 @@ def _prepare_image(self, image):
         dims = tf.shape(img)
         width = self.max_width
 
-        max_width = tf.to_int32(tf.ceil(tf.truediv(dims[1], dims[0]) * self.height_float))
-        max_height = tf.to_int32(tf.ceil(tf.truediv(width, max_width) * self.height_float))
+        max_width = tf.cast(tf.math.ceil(tf.truediv(dims[1], dims[0]) * self.height_float), dtype=tf.int32)
+        max_height = tf.cast(tf.math.ceil(tf.truediv(width, max_width) * self.height_float), dtype=tf.int32)
 
         resized = tf.cond(
             tf.greater_equal(width, max_width),
             lambda: tf.cond(
                 tf.less_equal(dims[0], self.height),
-                lambda: tf.to_float(img),
-                lambda: tf.image.resize_images(img, [self.height, max_width],
+                lambda: tf.cast(img, dtype=tf.float32),
+                lambda: tf.image.resize(img, [self.height, max_width],
                                                method=tf.image.ResizeMethod.BICUBIC),
             ),
-            lambda: tf.image.resize_images(img, [max_height, width],
+            lambda: tf.image.resize(img, [max_height, width],
                                            method=tf.image.ResizeMethod.BICUBIC)
         )