Speeding Up Tensorflow 2.0 Gradient Tape

I have been following the TF 2.0 tutorial for convolution VAE's, located here. Since it is eager, the gradients are computed by hand and then applied manually, using tf.GradientTap

Solution 1:

I found the solution: TensorFlow 2.0 introduces the concept of functions, which translate eager code into graph code.

The usage is pretty straight-forward. The only change needed is that all relevant functions (like compute_loss and apply_gradients) have to be annotated with @tf.function.

Solution 2:

The main reason why your code is taking a long time is because when we use a normal pythonic for loop which does not have any tensor, the construction of the graph takes tremendous amount of time because intuitively we might think that the same graph for training might be getting re used at every iteration, but the graph constructed is actually a chain like structure where each node is the training subgraph, and the total no. Of nodes in that chain are equivalent to the no. Of iterations in the loop. In a nutshell tensorflow unwraps the iterations and then constructs the graph. So it employs lot of redundancy in terms of both space and time for the graph alone. It's so bad that, just to add two tensors repeatedly in a normal pythonic loop for like a billion times, it takes almost half hour.

To get around this problem, specifically in your case we can take the help of .repeat transformation in the tf.data.Datasets api, instead of writing

for i in range(epochs) :

We can write

For x in x_train.repeat(epochs) :
        Train here