How can I list all Tensorflow variables a node depends on?

ghz 7months ago ⋅ 58 views

How can I list all Tensorflow variables/constants/placeholders a node depends on?

Example 1 (addition of constants):

import tensorflow as tf

a = tf.constant(1, name = 'a')
b = tf.constant(3, name = 'b')
c = tf.constant(9, name = 'c')
d = tf.add(a, b, name='d')
e = tf.add(d, c, name='e')

sess = tf.Session()
print(sess.run([d, e]))

I would like to have a function list_dependencies() such as:

  • list_dependencies(d) returns ['a', 'b']
  • list_dependencies(e) returns ['a', 'b', 'c']

Example 2 (matrix multiplication between a placeholder and a weight matrix, followed by the addition of a bias vector):

tf.set_random_seed(1)
input_size  = 5
output_size = 3
input       = tf.placeholder(tf.float32, shape=[1, input_size], name='input')
W           = tf.get_variable(
                "W",
                shape=[input_size, output_size],
                initializer=tf.contrib.layers.xavier_initializer())
b           = tf.get_variable(
                "b",
                shape=[output_size],
                initializer=tf.constant_initializer(2))
output      = tf.matmul(input, W, name="output")
output_bias = tf.nn.xw_plus_b(input, W, b, name="output_bias")

sess = tf.Session()
sess.run(tf.global_variables_initializer())
print(sess.run([output,output_bias], feed_dict={input: [[2]*input_size]}))

I would like to have a function list_dependencies() such as:

  • list_dependencies(output) returns ['W', 'input']
  • list_dependencies(output_bias) returns ['W', 'b', 'input']

Answers

To list all the TensorFlow variables, constants, and placeholders a node depends on, you can use TensorFlow's graph traversal utilities to examine the dependencies of each operation. Here's how you can achieve this:

  1. Obtain the list of dependencies by traversing the computational graph.
  2. Identify the type of each dependency (e.g., constant, variable, placeholder).

Here's the implementation for both TensorFlow 1.x and TensorFlow 2.x (with tf.compat.v1).

TensorFlow 1.x

import tensorflow as tf

def list_dependencies(op):
    deps = set()
    queue = [op]

    while queue:
        current_op = queue.pop(0)
        for inp in current_op.inputs:
            if inp.op.type in ["Placeholder", "VariableV2", "Const"]:
                deps.add(inp.op.name)
            queue.append(inp.op)
    
    return list(deps)

# Example 1
tf.reset_default_graph()
a = tf.constant(1, name='a')
b = tf.constant(3, name='b')
c = tf.constant(9, name='c')
d = tf.add(a, b, name='d')
e = tf.add(d, c, name='e')

sess = tf.Session()
print(sess.run([d, e]))
print("Dependencies of d:", list_dependencies(d.op))
print("Dependencies of e:", list_dependencies(e.op))

# Example 2
tf.reset_default_graph()
tf.set_random_seed(1)
input_size  = 5
output_size = 3
input       = tf.placeholder(tf.float32, shape=[1, input_size], name='input')
W           = tf.get_variable("W", shape=[input_size, output_size], initializer=tf.contrib.layers.xavier_initializer())
b           = tf.get_variable("b", shape=[output_size], initializer=tf.constant_initializer(2))
output      = tf.matmul(input, W, name="output")
output_bias = tf.nn.xw_plus_b(input, W, b, name="output_bias")

sess = tf.Session()
sess.run(tf.global_variables_initializer())
print(sess.run([output, output_bias], feed_dict={input: [[2]*input_size]}))
print("Dependencies of output:", list_dependencies(output.op))
print("Dependencies of output_bias:", list_dependencies(output_bias.op))

TensorFlow 2.x with tf.compat.v1

import tensorflow as tf

def list_dependencies(op):
    deps = set()
    queue = [op]

    while queue:
        current_op = queue.pop(0)
        for inp in current_op.inputs:
            if inp.op.type in ["Placeholder", "VariableV2", "Const"]:
                deps.add(inp.op.name)
            queue.append(inp.op)
    
    return list(deps)

# Example 1
tf.compat.v1.reset_default_graph()
a = tf.constant(1, name='a')
b = tf.constant(3, name='b')
c = tf.constant(9, name='c')
d = tf.add(a, b, name='d')
e = tf.add(d, c, name='e')

with tf.compat.v1.Session() as sess:
    print(sess.run([d, e]))
    print("Dependencies of d:", list_dependencies(d.op))
    print("Dependencies of e:", list_dependencies(e.op))

# Example 2
tf.compat.v1.reset_default_graph()
tf.compat.v1.set_random_seed(1)
input_size  = 5
output_size = 3
input       = tf.compat.v1.placeholder(tf.float32, shape=[1, input_size], name='input')
W           = tf.compat.v1.get_variable("W", shape=[input_size, output_size], initializer=tf.compat.v1.keras.initializers.VarianceScaling())
b           = tf.compat.v1.get_variable("b", shape=[output_size], initializer=tf.compat.v1.constant_initializer(2))
output      = tf.matmul(input, W, name="output")
output_bias = tf.compat.v1.nn.xw_plus_b(input, W, b, name="output_bias")

with tf.compat.v1.Session() as sess:
    sess.run(tf.compat.v1.global_variables_initializer())
    print(sess.run([output, output_bias], feed_dict={input: [[2]*input_size]}))
    print("Dependencies of output:", list_dependencies(output.op))
    print("Dependencies of output_bias:", list_dependencies(output_bias.op))

Explanation

  1. list_dependencies(op) Function:

    • This function takes a TensorFlow operation (op) and returns a list of dependencies (variables, constants, placeholders).
    • It uses a breadth-first search (BFS) to traverse the computational graph starting from the given operation.
    • For each operation, it checks its inputs and adds the names of placeholders, variables, and constants to the deps set.

  2. Examples:

    • For the addition of constants (Example 1), the function correctly identifies the dependencies (a, b, and c).
    • For the matrix multiplication and bias addition (Example 2), it correctly identifies the dependencies (input, W, and b).

This approach ensures that you can list all TensorFlow variables, constants, and placeholders a node depends on, regardless of the complexity of your graph.

#python #Tensorflow