Speeding Up Model Training with Google Colab

Accessing and making use of Google Colab’s GPUs.

Google Colab is a cloud computing service provided by Google, which can be utilized even by non paying users (whether this service is truly free or not is a different story). One of the highlights of Google Colab is the provision of hardware accelerators such as GPUs and even TPUs which can be used to train deep learning models at a much faster speed than on a CPU. This short article explains how to access and use the GPUs on Colab with either TensorFlow or PyTorch.

Speed your deep learning model training up using hardware accelerators. Photo by CHUTTERSNAP on Unsplash.

Setting Up the Hardware Accelerator on Colab

Before we even start writing any Python code, we need to first set up Colab’s runtime environment to use GPUs or TPUs instead of CPUs. Colab’s notebooks use CPUs by default — to change the runtime type to GPUs or TPUs, select “Change runtime type” under “Runtime” from Colab’s menu bar.

The hardware settings can be accessed from “Change runtime type” under “Runtime” in Colab’s menu bar. Image created by the author.

This brings up the notebook settings menu which allows you to choose the hardware accelerator. By default the hardware accelerator will be set to none, and you should be able to choose either the GPU or TPU from the drop down menu — for now select GPU.

Using the GPU with TensorFlow

In order to use the GPU with TensorFlow, obtain the device name using tf.test.gpu_device_name(). If the notebook is connected to a GPU, device_name will be set to /device:GPU:0. If not, we set device_name = /device:CPU:0 in order to use the CPU.

import tensorflow as tf
device_name = tf.test.gpu_device_name()
if len(device_name) > 0:
    print("Found GPU at: {}".format(device_name))
else:
    device_name = "/device:CPU:0"
    print("No GPU, using {}.".format(device_name))

Calculations to be performed on the GPU should be indented under a block specifying the device name. For example, to create and compile a TensorFlow model on the GPU, use the following outline below.

with tf.device(device_name):
    model = ...        # Create a TensorFlow model.
    model.compile(...) # Compile the model on the GPU.
# Once the model has been created and compiled on the GPU, it can be
# trained as per usual.
model.fit(X, y)

Using the GPU with PyTorch

The GPU device name can be obtained using PyTorch by checking if the GPU is available through torch.cuda.is_available(). If the GPU is available, we set device_name = torch.device(“cuda”), if not we use the CPU by setting device_name = torch.device(“cpu”).

import torch
if torch.cuda.is_available():
    device_name = torch.device("cuda")
else:
    device_name = torch.device('cpu')
print("Using {}.".format(device_name))

Unlike TensorFlow where the model is created and compiled on the GPU using a with indented code block, with PyTorch we can directly set the model to run on the GPU directly. However, in order for the PyTorch model to interact with the input data, both the model and the data must exist on the same device. Therefore, we need to set the data to run on the GPU as well.

model = ...           # Create a PyTorch model.
model.to(device_name) # Set the model to run on the GPU.
X = ...               # Load some data.
X = X.to(device_name) # Set the data to run on the GPU.
# Once the model and the data are set on the same device, 
# they can interact with each other. If not, a runtime error 
# will occur.
pred = model(X)     

References

1. https://www.tensorflow.org/guide/gpu
2. https://colab.research.google.com/github/phlippe/uvadlc_notebooks/blob/master/docs/tutorial_notebooks/tutorial2/Introduction_to_PyTorch.ipynb