Reusing variables with scoping

Until now, we have looked at the architecture of TensorFlow and the essentials required to implement a basic TensorFlow client. However, there is much more to TensorFlow than this. As we already saw, TensorFlow behaves quite differently from a typical Python script. For example, you cannot debug TensorFlow code in real time (as you would do a simple Python script using a Python IDE), as the computations do not happen in real time in TensorFlow (unless you are using the Eager Execution method, which was only recently in TensorFlow 1.7: https://research.googleblog.com/2017/10/eager-execution-imperative-define-by.html). In other words, TensorFlow first defines the full computational graph, does all computations on a device, and finally fetches results. Consequently, it can be quite tedious and painful to debug a TensorFlow client. This emphasizes the importance of attention to detail when implementing a TensorFlow client. Therefore, it is advised to adhere to proper coding practices introduced for TensorFlow. One such practice is known as scoping and allows easier variable reusing.

Reusing TensorFlow variables is a common scenario that occurs frequently in TensorFlow clients. To understand the value of an answer, we must first understand the question. Also, what better way to understand the question than erroneous code. Let's say that we want a function that performs a certain computation; given w, we need to compute x*w + y**2. Let's write a TensorFlow client, which has a function that performs this:

import tensorflow as tf
session = tf.InteractiveSession()
def very_simple_computation(w):
  x = tf.Variable(tf.constant(5.0, shape=None, dtype=tf.float32),
  name='x')
  y = tf.Variable(tf.constant(2.0, shape=None, dtype=tf.float32),
  name='y')
  z = x*w + y**2
  return z

Say that you want to compute this for a single step. Then, you can call session.run(very_simple_computation(2)) (of course, after calling tf.global_variables_initializer().run()), and you will have the answer and feel good about writing code that actually works. However, don't get too comfortable, because an issue arises if you want to run this function several times. Each time you call this method, two TensorFlow variables will be created. Remember that we discussed that TensorFlow is different to Python? This is one such instance. The x and y variables will not get replaced in the graph when you call this method multiple times. Rather, the old variables will be retained and new variables will be created in the graph until you run out of memory. But of course, the answer will be correct. To see this in action, run session.run(very_simple_computation(2)) in a for loop, and if you print the names of the variables in the graph, you will see more than two variables. This is the output when you run it 10 times:

'x:0', 'y:0', 'x_1:0', 'y_1:0', 'x_2:0', 'y_2:0', 'x_3:0', 'y_3:0', 'x_4:0', 'y_4:0', 'x_5:0', 'y_5:0', 'x_6:0', 'y_6:0', 'x_7:0', 'y_7:0', 'x_8:0', 'y_8:0', 'x_9:0', 'y_9:0', 'x_10:0', 'y_10:0'

Each time you run the function, a pair of variables is created. Let's make this explicit: if you run this function for 100 times, you will have 198 obsolete variables in your graph (99 x variables and 99 y variables).

This is where scoping comes to the rescue. Scoping allows you to reuse the variables instead of creating one each time a function is invoked. Now to add reusability to our little example, we will be changing the code to the following:

def not_so_simple_computation(w):
  x = tf.get_variable('x', initializer=tf.constant (5.0, shape=None, dtype=tf.float32))
  y = tf.get_variable('y', initializer=tf.constant(2.0, shape=None, dtype=tf.float32)) 
  z = x*w + y**2
  return z

def another_not_so_simple_computation(w):
  x = tf.get_variable('x', initializer=tf.constant(5.0, shape=None, dtype=tf.float32))
  y = tf.get_variable('y', initializer=tf.constant(2.0, shape=None, dtype=tf.float32)) 
  z = w*x*y
  return z
 
# Since this is the first call, the variables will # be created with following names
# x => scopeA/x, y => scopeA/y
with tf.variable_scope('scopeA'):
  z1 = not_so_simple_computation(tf.constant(1.0,dtype=tf.float32))
# scopeA/x and scopeA/y alread created we reuse them
with tf.variable_scope('scopeA',reuse=True):
  z2 = another_not_so_simple_computation(z1)

# Since this is the first call, the variables will be created with # be created with
# following names x => scopeB/x, y => scopeB/y
with tf.variable_scope('scopeB'):
  a1 = not_so_simple_computation(tf.constant(1.0,dtype=tf.float32))
# scopeB/x and scopeB/y alread created we reuse them
with tf.variable_scope('scopeB',reuse=True):
  a2 = another_not_so_simple_computation(a1)

# Say we want to reuse the "scopeA" again, since variables are already
# created we should set "reuse" argument to True when invoking the scope
with tf.variable_scope('scopeA',reuse=True):
  zz1 = not_so_simple_computation(tf.constant(1.0,dtype=tf.float32))
  zz2 = another_not_so_simple_computation(z1)

In this example, if you do session.run([z1,z2,a1,a2,zz1,zz2]), you should see z1, z2, a1, a2, zz1, zz2 has 9.0, 90.0, 9.0, 90.0, 9.0, 90.0 values in that order. Now if you print variables, you should see only four different variables: scopeA/x, scopeA/y, scopeB/x, and scopeB/y. We can now run it as many times as we want in a loop without worrying about creating redundant variables and running out of memory.

Now you might wonder why you cannot just create four variables at the beginning of the code and use them within the methods. However, this breaks the encapsulation of your code, because now you are explicitly depending on something outside your code.

Finally, scoping enables reusability while preserving the encapsulation of the code. Furthermore, scoping makes the flow of the code more intuitive and reduces the chance of errors as we are explicitly getting the variable by the scope and name instead of using the Python variable the TensorFlow variable was assigned to.