Updated for TensorFlow 2.x

ablation/ab_fixmatch_adam.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_anchoring.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_linearlr.py

@@ -14,7 +14,7 @@
 import functools
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_meanteacher.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_meanteacher_train.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_momentum.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_ra.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 from tqdm import trange

ablation/ab_fixmatch_sharpen.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

ablation/ab_fixmatch_sharpen_conditional_consistency.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

fixmatch.py

@@ -15,7 +15,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 from tqdm import trange

fully_supervised/fs_baseline.py

@@ -17,7 +17,7 @@
 import functools
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

fully_supervised/fs_mixup.py

@@ -16,7 +16,7 @@
 
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

fully_supervised/lib/train.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import flags
 from tqdm import trange
 

ict.py

@@ -19,7 +19,7 @@
 import functools
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

imagenet/augment/augment_ops.py

@@ -15,10 +15,9 @@
 
 import math
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 import tensorflow_addons as tfa
 
-
 # Default replace value
 REPLACE_VALUE = 128
 

imagenet/augment/ct_augment.py

@@ -13,11 +13,10 @@
 # limitations under the License.
 """Control theory augment."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from augment import augment_ops
 
-
 IMAGENET_AUG_OPS = [
  'AutoContrastBlend',
  'Blur',

imagenet/augment/rand_augment.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Random augment."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from augment import augment_ops
 

imagenet/augment/weak_augment.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Various weak agumentations."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def flip_augmentation(image):

imagenet/datasets/imagenet.py

@@ -14,12 +14,11 @@
 """Efficient ImageNet input pipeline using tf.data.Dataset."""
 
 import os
-from absl import flags
-
-import tensorflow as tf
 
+import tensorflow.compat.v1 as tf
+from absl import flags
 
-IMAGENET_BUFFER_SIZE = 16 * 1024 * 1024  # 16 MiB
+IMAGENET_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MiB
 IMAGENET_FETCH_CYCLE_LENGTH = 16
 IMAGENET_SHUFFLE_BUFFER_SIZE = 1024
 IMAGENET_PREPROCESSING_THREADS = 2

imagenet/fixmatch.py

@@ -13,21 +13,19 @@
 # limitations under the License.
 """Fixmatch semi-supervised training code."""
 
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 from absl import logging
 from easydict import EasyDict
 
-import tensorflow as tf
-
 import training
 from augment.augment import create_ssl_augmenter
 from datasets import datasets
 from models import resnet50_model
 from utils import ema
 from utils.learning_rate import ConfigurableLearningRateSchedule
 
-
 FLAGS = flags.FLAGS
 
 

imagenet/models/resnet50_model.py

@@ -19,7 +19,7 @@
 - http://torch.ch/blog/2016/02/04/resnets.html
 """
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 backend = tf.keras.backend

imagenet/prep_imagenet_data/save_sharded_data.py

@@ -17,15 +17,12 @@
 import os
 import time
 
+import apache_beam as beam
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
-
-import apache_beam as beam
 from apache_beam.options.pipeline_options import PipelineOptions
 
-import tensorflow as tf
-
-
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string('input_dir', None, 'Directory with input data.')

imagenet/supervised.py

@@ -13,20 +13,18 @@
 # limitations under the License.
 """Supervised training code."""
 
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 from absl import logging
 
-import tensorflow as tf
-
 import training
 from augment.augment import create_augmenter
 from datasets import datasets
 from models import resnet50_model
 from utils import ema
 from utils.learning_rate import ConfigurableLearningRateSchedule
 
-
 FLAGS = flags.FLAGS
 
 

imagenet/training.py

@@ -20,13 +20,11 @@
 import os
 import time
 
+import tensorflow.compat.v1 as tf
 from absl import flags
 from absl import logging
 from easydict import EasyDict
 
-import tensorflow as tf
-
-
 flags.DEFINE_string(
  'tpu',
  None,

imagenet/utils/ema.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Helper functions for exponential moving average."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 def assign_ema_vars_from_initial_values(ema_variables, initial_values):

imagenet/utils/learning_rate.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Utils for learning rate schedule."""
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 
 class ConfigurableLearningRateSchedule(

libml/augment.py

@@ -20,7 +20,7 @@
 import random
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import flags
 
 from libml import utils, ctaugment

libml/data.py

@@ -19,7 +19,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 from tqdm import tqdm

libml/layers.py

@@ -16,7 +16,7 @@
 import numbers
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 
 from libml.data import DataSets
 

libml/models.py

@@ -15,7 +15,7 @@
 import functools
 import itertools
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import flags
 
 from libml import layers

libml/train.py

@@ -18,7 +18,7 @@
 import shutil
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import flags
 from tqdm import trange, tqdm
 

libml/utils.py

@@ -17,7 +17,7 @@
 import re
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import flags, logging
 from tensorflow.python.client import device_lib
 
@@ -49,6 +49,7 @@ def setup_main():
 def setup_tf():
  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
  logging.set_verbosity(logging.ERROR)
+ tf.disable_v2_behavior()
 
 
 def smart_shape(x):

mean_teacher.py

@@ -21,7 +21,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

mixmatch.py

@@ -22,7 +22,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

mixup.py

@@ -18,7 +18,7 @@
 import functools
 import os
 
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

pi_model.py

@@ -20,7 +20,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags
 

pseudo_label.py

@@ -20,7 +20,7 @@
 import os
 
 import numpy as np
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from absl import app
 from absl import flags