Tensorflow 2.0: How can I fully customize a Tensorflow training loop like I can with PyTorch?
Tensorflow uses eager execution, it means your graph is connected dynamically. Offcourse it’s not as dynamic as PyTorch. But, google is trying its best and incorporated a lot of features in 2.0 and beyond.
For writing a custom loop in Tensorflow you need to use tf.GradientTape. I would say it involves the same steps as in pytorch. Compute gradients and then update using optimizer.
I used to use Tensorflow a lot before, but moved over to Pytorch because it was just a lot easier to debug. The nice thing I found with PyTorch is that I have to write my own training loop, so I can step through the code and find errors. I can fire up pdb and check the tensor shapes and transformations, etc., without difficulty.
In Tensorflow I was using the model.fit() function all the time, and so any error message I got was like 6 pages of C code where the error message did not give me any indication was in the python code. User's can't step through the model.fit() function since it is a static graph, so that really slowed down my development process. BUT, I was thinking about using Tensorflow again and I was wondering whether a user can step through a custom training loop and look at tensor shapes, etc., or whether even a custom training loop is compiled to a static graph and hence users cannot step through it?
I did google this question, but all of the tutorials for custom training loops in Tensorflow focus on custom loops being for advanced users, such as if you want to apply some exotic callback while training or if you want to apply some conditional logic. So the simple question of whether it is easy to step through a custom training loop is not answered.
Any help is appreciated. Thanks.