When I trained a model with model.fit with my own data generator function in tensorflow keras, I noticed GPU utilization could be almost 100%. This is because the model I trained is a good one regarding the parallelization.
Nonetheless if I built a train function, something like the following function, I noticed the utilization is very different:
@tf.function
def train_step(input, output):
with tf.GradientTape() as tape:
predictions = model(input)
loss = loss_object(output, predictions)
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_loss(loss)
In detail, I noticed it goes from 0% to almost 100%, and then goes down to 0% again for a moment between updates. I guess between two updates, data needs to transfer from CPU to GPU in the first place. This explains why it goes from 100% to 0% for some moment.
But this leads to a question to me is: How to make it more optimal for a custom training function with tf.GradientTape so that the GPUs is almost 100% all the time (as with model.fit)?