How to evaluate only a random subset of all possible operations per pass inside a graph?

user_8690766 0
Viewed 47
import tensorflow as tf

operations = [
    tf.keras.layers.Dense(64, activation='sigmoid'),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation=None),
]


def call(x):
    # Sample 2 of the 3 operations.
    sampled_ids = tf.random.categorical(
        tf.zeros((1, len(operations))), num_samples=2, dtype=tf.int32,
    )[0]

    # Compute their output given x.
    op_results_eager = tf.stack([operations[op_id](x) for op_id in sampled_ids])

    # Try to replicate op_results_eager in graph_mode without evaluating all operations!
    op_result_functs = [lambda: op(x) for op in operations]
    op_results_graph = tf.stack([tf.switch_case(branch_index=op_id, branch_fns=op_result_functs) for op_id in sampled_ids])

    tf.print(tf.reduce_all(tf.equal(op_results_eager, op_results_graph)))

    return op_results_graph


for _ in range(1000):
    call(tf.ones(shape=(1, 5)))

I need the same result as in op_results_eager but with statements that allow wrapping call() as tf.function while still evaluating only the sampled operations.

As you can see, I tried to build the indexing as a switch case, but that doesn't even give the right result in eager execution.

tensorflow python
1 Answers

Maybe something like this (although each Dense layer receives data each time):

import tensorflow as tf

dense1 = tf.keras.layers.Dense(64, activation='sigmoid')
dense2 = tf.keras.layers.Dense(64, activation='relu')
dense3 = tf.keras.layers.Dense(64, activation=None)

operations = [dense1, dense2, dense3]

def call(x):
    # Sample 2 of the 3 operations.
    sampled_ids = tf.random.categorical(
        tf.zeros((1, len(operations))), num_samples=2, dtype=tf.int32,
    )[0]

    op_results_eager = tf.stack([operations[op_id](x) for op_id in sampled_ids])
 
    op_results_graph = tf.gather(tf.stack([dense1(x), dense2(x), dense3(x)]), sampled_ids)

    tf.print(tf.reduce_all(tf.equal(op_results_eager, op_results_graph)))

    return op_results_graph


for _ in range(10):
    call(tf.ones(shape=(1, 5)))

1
1
1
1
1
1
1
1
1
1

Or another alternative without calling each layer:

dense1 = tf.keras.layers.Dense(64, activation='sigmoid')
dense2 = tf.keras.layers.Dense(64, activation='relu')
dense3 = tf.keras.layers.Dense(64, activation=None)

operations = [dense1, dense2, dense3]

def call(x):
    # Sample 2 of the 3 operations.
    sampled_ids = tf.random.categorical(
        tf.zeros((1, len(operations))), num_samples=2, dtype=tf.int32,
    )[0]

    op_results_eager = tf.stack([operations[op_id](x) for op_id in sampled_ids])
    
    id1, id2 = tf.split(sampled_ids, 2)
  
    output1 = tf.stack([tf.cond(tf.equal(id1, 0), lambda: dense1(x), lambda: tf.zeros((1, 64))),tf.cond(tf.equal(id2, 0), lambda: dense1(x), lambda: tf.zeros((1, 64)))])
    output2 = tf.stack([tf.cond(tf.equal(id1, 1), lambda: dense2(x), lambda: tf.zeros((1, 64))),tf.cond(tf.equal(id2, 1), lambda: dense2(x), lambda: tf.zeros((1, 64)))])
    output3 = tf.stack([tf.cond(tf.equal(id1, 2), lambda: dense3(x), lambda: tf.zeros((1, 64))),tf.cond(tf.equal(id2, 2), lambda: dense3(x), lambda: tf.zeros((1, 64)))])
    outputs = tf.stack([output1, output2, output3], axis=0)

    op_results_graph = tf.expand_dims(tf.gather_nd(outputs, tf.where(tf.reduce_any(tf.not_equal(outputs, tf.zeros((64,))), axis=-1))), axis=1)

    return op_results_graph
user_96578610
Related