Question about GPU usage in Google Colab when training Keras/TF models

Question

I have a quick question: when using Google Colab with the GPU enabled, does all of the code already run on the GPU then or is there some setting in the code that we must change to make it run on the GPU? Specifically, if I am training a neural network model in Keras/TF, do I need to edit my code in any way to ensure that the model is trained on the GPU?

Thanks!

Answer 1

as showed in the Tensorflow with Gpu example notebook you can run your model in the following way to make sure it is running on the chosen device:

def gpu():
  with tf.device('/device:GPU:0'):
  random_image_gpu = tf.random.normal((100, 100, 100, 3))
  net_gpu = tf.keras.layers.Conv2D(32, 7)(random_image_gpu)
  return tf.math.reduce_sum(net_gpu)