Merge pull request #3 from ercaronte/feature/tf2compatibility

Upgrade lucid to use tensorflow v2 APIs and make it work on Colab platform

Author: ercaronte

lucid/__init__.py

@@ -35,4 +35,4 @@
 seed = 0
 
 # Set the lucid version - setup.py imports this value!
-__version__ = "0.3.10"
+__version__ = "0.3.11"

lucid/misc/gl/glcontext.py

@@ -37,7 +37,7 @@
         'prior importing this module.')
   raise
 
-import ctypes
+import ctypes.util
 from ctypes import pointer
 import os
 
@@ -117,4 +117,4 @@ def create_opengl_context(surface_size=(640, 480)):
 
   egl_context = egl.eglCreateContext(egl_display, egl_cfg, egl.EGL_NO_CONTEXT,
                                      None)
-  egl.eglMakeCurrent(egl_display, egl_surf, egl_surf, egl_context)
+  egl.eglMakeCurrent(egl_display, egl_surf, egl_surf, egl_context)

lucid/misc/gradient_override.py

@@ -100,7 +100,7 @@ def _foo_grad_alt(op, grad): ...
        override_dict_by_name[op_name] = grad_f
     else:
       override_dict_by_name[op_name] = register_to_random_name(grad_f)
-  with tf.get_default_graph().gradient_override_map(override_dict_by_name):
+  with tf.compat.v1.get_default_graph().gradient_override_map(override_dict_by_name):
     yield
 
 
@@ -154,7 +154,8 @@ def store_out(out_value):
         state["out_value"] = out_value
 
       store_name = "store_" + f.__name__
-      store = tf.py_func(store_out, [out], (), stateful=True, name=store_name)
+      store = tf.compat.v1.py_func(store_out, [out], (), stateful=True, name=store_name)
+      # store = tf.numpy_function(store_out, [out], (), stateful=True, name=store_name) # not yet implemented in TF 2.5
 
       # Next, we create the mock function, with an overriden gradient.
       # Note that we need to make sure store gets evaluated before the mock
@@ -167,8 +168,8 @@ def mock_f(*inputs):
       with tf.control_dependencies([store]):
         with gradient_override_map({"PyFunc": grad_f_name}):
           mock_name = "mock_" + f.__name__
-          mock_out = tf.py_func(mock_f, inputs, out.dtype, stateful=True,
-                                name=mock_name)
+          mock_out = tf.compat.v1.py_func(mock_f, inputs, out.dtype, stateful=True, name=mock_name)
+          # mock_out = tf.numpy_function(mock_f, inputs, out.dtype, stateful=True, name=mock_name) # not yet implemented in TF 2.5
           mock_out.set_shape(out.get_shape())
 
       # Finally, we can return the mock.

lucid/misc/io/loading.py

@@ -126,7 +126,7 @@ def _load_text(handle, split=False, encoding="utf-8"):
 def _load_graphdef_protobuf(handle, **kwargs):
     """Load GraphDef from a binary proto file."""
     # as_graph_def
-    graph_def = tf.GraphDef.FromString(handle.read())
+    graph_def = tf.compat.v1.GraphDef.FromString(handle.read())
 
     # check if this is a lucid-saved model
     # metadata = modelzoo.util.extract_metadata(graph_def)

lucid/misc/io/reading.py

@@ -25,7 +25,7 @@
 import logging
 from urllib.parse import urlparse
 from urllib import request
-from tensorflow.io.gfile import GFile
+from tensorflow.compat.v1.io.gfile import GFile
 import tensorflow as tf
 from tempfile import gettempdir
 import gc

lucid/misc/io/showing.py

@@ -329,7 +329,7 @@ def _strip_consts(graph_def, max_const_size=32):
     This is mostly a utility function for graph(), and also originates here:
     https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/deepdream/deepdream.ipynb
     """
-    strip_def = tf.GraphDef()
+    strip_def = tf.compat.v1.GraphDef()
     for n0 in graph_def.node:
         n = strip_def.node.add()
         n.MergeFrom(n0)

lucid/misc/io/writing.py

@@ -25,7 +25,7 @@
 import logging
 from contextlib import contextmanager
 from urllib.parse import urlparse
-from tensorflow import gfile
+from tensorflow.compat.v1.io import gfile
 
 from lucid.misc.io.scoping import scope_url
 
@@ -66,14 +66,14 @@ def write_handle(path, mode=None):
     path = scope_url(path)
 
     if _supports_make_dirs(path):
-        gfile.MakeDirs(os.path.dirname(path))
+        gfile.makedirs(os.path.dirname(path))
 
     if mode is None:
         if _supports_binary_writing(path):
             mode = "wb"
         else:
             mode = "w"
 
-    handle = gfile.Open(path, mode)
+    handle = gfile.GFile(path, mode)
     yield handle
     handle.close()

lucid/misc/redirected_relu_grad.py

@@ -92,12 +92,12 @@ def redirected_relu_grad(op, grad):
   x = op.inputs[0]
 
   # Compute ReLu gradient
-  relu_grad = tf.where(x < 0., tf.zeros_like(grad), grad)
+  relu_grad = tf.compat.v1.where(x < 0., tf.zeros_like(grad), grad)
 
   # Compute redirected gradient: where do we need to zero out incoming gradient
   # to prevent input going lower if its already negative
   neg_pushing_lower = tf.logical_and(x < 0., grad > 0.)
-  redirected_grad = tf.where(neg_pushing_lower, tf.zeros_like(grad), grad)
+  redirected_grad = tf.compat.v1.where(neg_pushing_lower, tf.zeros_like(grad), grad)
 
   # Ensure we have at least a rank 2 tensor, as we expect a batch dimension
   assert_op = tf.Assert(tf.greater(tf.rank(relu_grad), 1), [tf.rank(relu_grad)])
@@ -106,12 +106,12 @@ def redirected_relu_grad(op, grad):
     batch = tf.shape(relu_grad)[0]
     reshaped_relu_grad = tf.reshape(relu_grad, [batch, -1])
     relu_grad_mag = tf.norm(reshaped_relu_grad, axis=1)
-  result_grad = tf.where(relu_grad_mag > 0., relu_grad, redirected_grad)
+  result_grad = tf.compat.v1.where(relu_grad_mag > 0., relu_grad, redirected_grad)
 
-  global_step_t =tf.train.get_or_create_global_step()
+  global_step_t = tf.compat.v1.train.get_or_create_global_step()
   return_relu_grad = tf.greater(global_step_t, tf.constant(16, tf.int64))
 
-  return tf.where(return_relu_grad, relu_grad, result_grad)
+  return tf.compat.v1.where(return_relu_grad, relu_grad, result_grad)
 
 
 def redirected_relu6_grad(op, grad):
@@ -120,15 +120,15 @@ def redirected_relu6_grad(op, grad):
 
   # Compute ReLu gradient
   relu6_cond = tf.logical_or(x < 0., x > 6.)
-  relu_grad = tf.where(relu6_cond, tf.zeros_like(grad), grad)
+  relu_grad = tf.compat.v1.where(relu6_cond, tf.zeros_like(grad), grad)
 
   # Compute redirected gradient: where do we need to zero out incoming gradient
   # to prevent input going lower if its already negative, or going higher if
   # already bigger than 6?
   neg_pushing_lower = tf.logical_and(x < 0., grad > 0.)
   pos_pushing_higher = tf.logical_and(x > 6., grad < 0.)
   dir_filter = tf.logical_or(neg_pushing_lower, pos_pushing_higher)
-  redirected_grad = tf.where(dir_filter, tf.zeros_like(grad), grad)
+  redirected_grad = tf.compat.v1.where(dir_filter, tf.zeros_like(grad), grad)
 
   # Ensure we have at least a rank 2 tensor, as we expect a batch dimension
   assert_op = tf.Assert(tf.greater(tf.rank(relu_grad), 1), [tf.rank(relu_grad)])
@@ -137,9 +137,9 @@ def redirected_relu6_grad(op, grad):
     batch = tf.shape(relu_grad)[0]
     reshaped_relu_grad = tf.reshape(relu_grad, [batch, -1])
     relu_grad_mag = tf.norm(reshaped_relu_grad, axis=1)
-  result_grad =  tf.where(relu_grad_mag > 0., relu_grad, redirected_grad)
+  result_grad =  tf.compat.v1.where(relu_grad_mag > 0., relu_grad, redirected_grad)
 
-  global_step_t = tf.train.get_or_create_global_step()
+  global_step_t = tf.compat.v1.train.get_or_create_global_step()
   return_relu_grad = tf.greater(global_step_t, tf.constant(16, tf.int64))
 
-  return tf.where(return_relu_grad, relu_grad, result_grad)
+  return tf.compat.v1.where(return_relu_grad, relu_grad, result_grad)

lucid/misc/tfutil.py

@@ -25,7 +25,7 @@ def create_session(target='', timeout_sec=10):
   when having multiple python sessions sharing the same GPU.
   '''
   graph = tf.Graph()
-  config = tf.ConfigProto()
+  config = tf.compat.v1.ConfigProto()
   config.gpu_options.allow_growth = True
   config.operation_timeout_in_ms = int(timeout_sec*1000)
-  return tf.InteractiveSession(target=target, graph=graph, config=config)
+  return tf.compat.v1.InteractiveSession(target=target, graph=graph, config=config)

lucid/modelzoo/other_models/Clip.py

@@ -21,3 +21,4 @@ class Clip_ResNet50_4x(Model):
     model_name = "Clip_ResNet50_4x"
     image_shape = [288, 288, 3]
     model_path = "gs://modelzoo/vision/other_models/Clip_ResNet50_4x.pb"
+    dataset = None

lucid/modelzoo/other_models/InceptionV1.py

@@ -21,7 +21,7 @@
 
 def _populate_inception_bottlenecks(scope):
   """Add Inception bottlenecks and their pre-Relu versions to the graph."""
-  graph = tf.get_default_graph()
+  graph = tf.compat.v1.get_default_graph()
   for op in graph.get_operations():
     if op.name.startswith(scope+'/') and 'Concat' in op.type:
       name = op.name.split('/')[1]
@@ -37,7 +37,7 @@ def _populate_inception_bottlenecks(scope):
 class InceptionV1(Model):
   """InceptionV1 (or 'GoogLeNet')
 
-  This is a (re?)implementation of InceptionV1 from the "Going deeper 
+  This is a (re?)implementation of InceptionV1 from the "Going deeper
   with convolutions" paper. Links:
     * Official CVPR paper, requires subscription: https://ieeexplore.ieee.org/document/7298594
     * Author preprint: https://www.cs.unc.edu/~wliu/papers/GoogLeNet.pdf

lucid/modelzoo/util.py

@@ -55,31 +55,31 @@ def forget_xy(t):
   filter) when we only use early parts of it.
   """
   shape = (t.shape[0], None, None, t.shape[3])
-  return tf.placeholder_with_default(t, shape)
+  return tf.compat.v1.placeholder_with_default(t, shape)
 
 
 def frozen_default_graph_def(input_node_names, output_node_names):
   """Return frozen and simplified graph_def of default graph."""
 
-  sess = tf.get_default_session()
+  sess = tf.compat.v1.get_default_session()
   if sess is None:
     raise RuntimeError("Default session not registered.")
 
-  input_graph_def = tf.get_default_graph().as_graph_def()
+  input_graph_def = tf.compat.v1.get_default_graph().as_graph_def()
   if len(input_graph_def.node) == 0:
     raise RuntimeError("Default graph is empty. Is it possible your model wasn't constructed or is in a different graph?")
 
-  pruned_graph = tf.graph_util.remove_training_nodes(
+  pruned_graph = tf.compat.v1.graph_util.remove_training_nodes(
       input_graph_def, protected_nodes=(output_node_names + input_node_names)
   )
-  pruned_graph = tf.graph_util.extract_sub_graph(pruned_graph, output_node_names)
+  pruned_graph = tf.compat.v1.graph_util.extract_sub_graph(pruned_graph, output_node_names)
 
   # remove explicit device assignments
   for node in pruned_graph.node:
       node.device = ""
 
-  all_variable_names = [v.op.name for v in tf.global_variables()]
-  output_graph_def = tf.graph_util.convert_variables_to_constants(
+  all_variable_names = [v.op.name for v in tf.compat.v1.global_variables()]
+  output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
       sess=sess,
       input_graph_def=pruned_graph,
       output_node_names=output_node_names,

lucid/modelzoo/vision_base.py

@@ -175,7 +175,7 @@ def post_import(self, scope):
   def create_input(self, t_input=None, forget_xy_shape=True):
     """Create input tensor."""
     if t_input is None:
-      t_input = tf.placeholder(tf.float32, self.image_shape)
+      t_input = tf.compat.v1.placeholder(tf.float32, self.image_shape)
     t_prep_input = t_input
     if len(t_prep_input.shape) == 3:
       t_prep_input = tf.expand_dims(t_prep_input, 0)
@@ -189,7 +189,7 @@ def create_input(self, t_input=None, forget_xy_shape=True):
 
   def import_graph(self, t_input=None, scope='import', forget_xy_shape=True, input_map=None):
     """Import model GraphDef into the current graph."""
-    graph = tf.get_default_graph()
+    graph = tf.compat.v1.get_default_graph()
     assert graph.unique_name(scope, False) == scope, (
         'Scope "%s" already exists. Provide explicit scope names when '
         'importing multiple instances of the model.') % scope
@@ -231,7 +231,7 @@ def get_layer(self, name):
   @staticmethod
   def suggest_save_args(graph_def=None):
     if graph_def is None:
-      graph_def = tf.get_default_graph().as_graph_def()
+      graph_def = tf.compat.v1.get_default_graph().as_graph_def()
     gdhelper = model_util.GraphDefHelper(graph_def)
     inferred_info = dict.fromkeys(("input_name", "image_shape", "output_names", "image_value_range"))
     node_shape = lambda n: [dim.size for dim in n.attr['shape'].shape.dim]

lucid/optvis/param/cppn.py

@@ -18,7 +18,8 @@
 
 import numpy as np
 import tensorflow as tf
-from tensorflow.contrib import slim
+# from tensorflow.contrib import slim
+import tf_slim as slim
 
 
 def _composite_activation(x, biased=True):
@@ -83,8 +84,8 @@ def cppn(
         [slim.conv2d],
         kernel_size=[1, 1],
         activation_fn=None,
-        weights_initializer=tf.initializers.variance_scaling(),
-        biases_initializer=tf.initializers.random_normal(0.0, 0.1),
+        weights_initializer=tf.compat.v1.initializers.variance_scaling(),
+        biases_initializer=tf.compat.v1.initializers.random_normal(0.0, 0.1),
     ):
         for i in range(num_layers):
             x = slim.conv2d(net, num_hidden_channels)

lucid/optvis/param/random.py

@@ -34,7 +34,7 @@ def rand_fft_image(shape, sd=None, decay_power=1):
     for _ in range(b):
         freqs = rfft2d_freqs(h, w)
         fh, fw = freqs.shape
-        spectrum_var = sd * tf.random_normal([2, ch, fh, fw], dtype="float32")
+        spectrum_var = sd * tf.random.normal([2, ch, fh, fw], dtype="float32")
         spectrum = tf.complex(spectrum_var[0], spectrum_var[1])
         spertum_scale = 1.0 / np.maximum(freqs, 1.0 / max(h, w)) ** decay_power
         # Scale the spectrum by the square-root of the number of pixels

lucid/optvis/param/resize_bilinear_nd.py

@@ -105,7 +105,7 @@ def resize_bilinear_nd(t, target_shape):
     # We can then reshape and use the 2d tf.image.resize_bilinear() on the
     # inner two dimesions.
     t_ = tf.reshape(t, shape_)
-    t_ = tf.image.resize_bilinear(t_, new_shape_[1:3])
+    t_ = tf.compat.v1.image.resize_bilinear(t_, new_shape_[1:3])
 
     # And then reshape back to our uncollapsed version, having finished resizing
     # two more dimensions in our shape.

lucid/optvis/param/spatial.py

@@ -125,11 +125,11 @@ def bilinearly_sampled_image(texture, uv):
     h, w = tf.unstack(tf.shape(texture)[:2])
     u, v = tf.split(uv, 2, axis=-1)
     v = 1.0 - v  # vertical flip to match GL convention
-    u, v = u * tf.to_float(w) - 0.5, v * tf.to_float(h) - 0.5
-    u0, u1 = tf.floor(u), tf.ceil(u)
-    v0, v1 = tf.floor(v), tf.ceil(v)
+    u, v = u * tf.cast(w, tf.float32) - 0.5, v * tf.cast(h, tf.float32) - 0.5
+    u0, u1 = tf.math.floor(u), tf.math.ceil(u)
+    v0, v1 = tf.math.floor(v), tf.math.ceil(v)
     uf, vf = u - u0, v - v0
-    u0, u1, v0, v1 = map(tf.to_int32, [u0, u1, v0, v1])
+    u0, u1, v0, v1 = map(lambda x: tf.cast(x, tf.int32), [u0, u1, v0, v1])
 
     def sample(u, v):
         vu = tf.concat([v % h, u % w], axis=-1)

lucid/optvis/render.py

@@ -86,16 +86,17 @@ def render_vis(model, objective_f, param_f=None, optimizer=None,
     multiple channel visualizations stacked on top of each other.
   """
 
-  with tf.Graph().as_default() as graph, tf.Session() as sess:
+  with tf.Graph().as_default() as graph, tf.compat.v1.Session() as sess:
+    # print(f'render vis graph: {graph}')
 
     if use_fixed_seed:  # does not mean results are reproducible, see Args doc
-      tf.set_random_seed(0)
+        tf.compat.v1.set_random_seed(0)
 
     T = make_vis_T(model, objective_f, param_f, optimizer, transforms,
                    relu_gradient_override)
     print_objective_func = make_print_objective_func(print_objectives, T)
     loss, vis_op, t_image = T("loss"), T("vis_op"), T("input")
-    tf.global_variables_initializer().run()
+    tf.compat.v1.global_variables_initializer().run()
 
     images = []
     try:
@@ -167,8 +168,8 @@ def make_vis_T(model, objective_f, param_f=None, optimizer=None,
   transform_f = make_transform_f(transforms)
   optimizer = make_optimizer(optimizer, [])
 
-  global_step = tf.train.get_or_create_global_step()
-  init_global_step = tf.variables_initializer([global_step])
+  global_step = tf.compat.v1.train.get_or_create_global_step()
+  init_global_step = tf.compat.v1.variables_initializer([global_step])
   init_global_step.run()
 
   if relu_gradient_override:
@@ -221,7 +222,7 @@ def make_t_image(param_f):
   if not isinstance(t_image, tf.Tensor):
     raise TypeError("param_f should produce a Tensor, but instead created a "
                    + str(type(t_image)) )
-  elif t_image.graph != tf.get_default_graph():
+  elif t_image.graph != tf.compat.v1.get_default_graph():
     raise TypeError("""param_f produced a t_image tensor belonging to a graph
                      that isn't the default graph for rendering. Did you
                      accidentally use render_vis when you meant to use
@@ -239,15 +240,15 @@ def make_transform_f(transforms):
 
 def make_optimizer(optimizer, args):
   if optimizer is None:
-    return tf.train.AdamOptimizer(0.05)
+    return tf.compat.v1.train.AdamOptimizer(0.05)
   elif callable(optimizer):
     return optimizer(*args)
-  elif isinstance(optimizer, tf.train.Optimizer):
+  elif isinstance(optimizer, tf.compat.v1.train.Optimizer):
     return optimizer
   else:
     raise ("Could not convert optimizer argument to usable optimizer. "
            "Needs to be one of None, function from (graph, sess) to "
-           "optimizer, or tf.train.Optimizer instance.")
+           "optimizer, or tf.compat.v1.train.Optimizer instance.")
 
 
 def import_model(model, t_image, t_image_raw=None, scope="import", input_map=None):