TensorBoard tf.cond performance

user_15550096 0
Viewed 158

I have written a very simple custom layeras follows:

class Custom_Layer1234(keras.layers.Layer):
    def __init__(self,  inputname , units=45, input_dim=45):
        super(Custom_Layer1234, self).__init__()
        w_init = tf.random_normal_initializer()
        b_init = tf.zeros_initializer()

        self.w_0 = tf.Variable(initial_value=w_init(shape=(input_dim, units,)),
                              name='w0{}'.format(inputname),  trainable=True,)
        self.w_1 = tf.Variable(initial_value=w_init(shape=(input_dim, units,)),
                              name='w1{}'.format(inputname),  trainable=True,)
        self.b_0 = tf.Variable(initial_value=b_init(shape=(units,), dtype="float32"),
                              name='b0{}'.format(inputname), trainable=True)
        self.b_1 = tf.Variable(initial_value=b_init(shape=(units,), dtype="float32"),
                              name='b1{}'.format(inputname), trainable=True)
    @tf.function    
    def call(self, inputs):
        diff_1 = inputs[0][1]
       
        if diff_1 <= 0 :    
            y = tf.matmul(inputs, self.w_0) + self.b_0
        else:
            y = tf.matmul(inputs, self.w_1) + self.b_1
           
        return tf.nn.relu(y)

I then sliced the input Tensor into smaller ones, looping it and feed each 'slice' into one of this custom layer, so I suspect that at runtime there is a variety of graphs of 'different branches' base on each input[i]. This model runs quite slowly, while GPU utilization is below 30% most of the time.

I opened Tensorboard to check how I can train the model faster and followed its instruction. my input pipeline has no issue. and the pre-processing is fine too. I also tried Tensorboard's suggestion on :

TF_GPU_THREAD_MODE=gpu_private
mixed-precision

and none cut the training time. I removed the if-condition in the layer and reran the model:

# if diff_1 <= 0 :   
   y = tf.matmul(inputs, self.w_0) + self.b_0
# else:
#  y = tf.matmul(inputs, self.w_1) + self.b_1

and is significantly faster, so I'm guessing that many of the tf.cond causes this slow down. I was thinking if there's a way to save these ' routes ' as many unique graphs (tf.Graph)individually for each input[i] and can be reused such that this computation doesn't need to be done repeatedly(is my assumption correct?). or TensorFlow already is doing this. Could I benefit from any change to the code or the mode to run, or the graph optimizer Grappler ? such that the training could be made faster.

Thank you very much

tensorboard keras tensorflow python
1 Answers

Introducing conditions in computational graph indeed significantly slows down TF performance. For example, using RNN layer with mask parameter significantly slows down its performance.

I am not an expert on this, but my understanding of this issue is the following. To update the parameters of a network, TF must trace all gradients. When there are no conditional statements, the same formula can be applied for all samples in the batch. On the contrary, with conditional statements the individual "paths" of the data through the computational graph must be known. My guess is that this makes vectorization (use of batches) ineffective.

For a similar reason, vectorization can become not really effective for the forward propagation of data through the network as well.

Considering code optimization, you may try the following. Make two distinct layers, layer_A and layer_B, that will do the calculations you want based on diff_1 value. Then, instead of using if statement, try

    # you may need to change indecies of inputs
    result = tf.where(inputs[:,0,1]>0, layer_A(inputs), layer_B(inputs))

Not sure that this will be faster, but I would give this a try.

user_127816740
Related