How To Use GPU with PyTorch
A short tutorial on using GPUs for your deep learning models with PyTorch. Made by Ayush Thakur using Weights & Biases
In this report, we will walk through ways to use and have more control over your GPU.
We'll use Weights and Biases that lets us automatically log all our GPU and CPU utilization metrics. This makes it easy to monitor the compute resource usage as we train a plethora of models.
Check GPU Availability
The easiest way to check if you have access to GPUs is to call torch.cuda.is_available(). If it returns True, it means the system has the Nvidia driver correctly installed.
>>> import torch>>> torch.cuda.is_available()
Use GPU - Gotchas
By default, the tensors are generated on the CPU. Even the model is initialized on the CPU. Thus one has to manually ensure that the operations are done using GPU.
>>> X_train = torch.FloatTensor([0., 1., 2.])>>> X_train.is_cudaFalse
PyTorch provides a simple to use API to transfer the tensor generated on CPU to GPU. Luckily the new tensors are generated on the same device as the parent tensor.
>>> X_train = X_train.to(device)>>> X_train.is_cudaTrue
The same logic applies to the model.
model = MyModel(args) model.to(device)
Thus data and the model need to be transferred to the GPU. Well, what's device?
It's a common PyTorch practice to initialize a variable, usually named device that will hold the device we’re training on (CPU or GPU).
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")print(device)
Torch CUDA Package
In PyTorch, the torch.cuda package has additional support for CUDA tensor types, that implement the same function as CPU tensors, but they utilize GPUs for computation.
If you want a tensor to be on GPU you can call .cuda().
>>> X_train = torch.FloatTensor([0., 1., 2.])>>> X_train = X_train.cuda()
If you have a tensor on GPU and you would like to bring it to CPU then you can call .cpu(). This is usually used to bring the output(tensor) of the model to the CPU.
To get the index of the currently selected device.
>>> torch.cuda.current_device() # returns 0 in my case
To get the number of GPUs available.
>>> torch.cuda.device_count() # returns 1 in my case
To get the name of the device.
>>> torch.cuda.get_device_name(0) # good old Tesla K80
Example and GPU Metrics Visualization
Try out the linked colab notebook to train a simple MNIST classifier using PyTorch. The notebook is integrated with Weights and Biases
If you are tracking your models using Weights & Biases, all your system metrics, including GPU utilization, will be automatically logged. Some of the most important metrics logged are GPU memory allocated, GPU utilization, CPU utilization, etc. You can see the full list of metrics logged here
The media panel shown below shows some of these system metrics that were automatically logged by W&B while training.
In this article, you saw how you can leverage GPUs for your deep learning research using Keras, and use Weights and Biases to monitor your resource consumption. Check out this great article by Lambda Labs on tracking system resource utilization
during training with the Weights & Biases.
Weights & Biases
Weights & Biases helps you keep track of your machine learning experiments. Use our tool to log hyperparameters and output metrics from your runs, then visualize and compare results and quickly share findings with your colleagues.
Recommended Reading For Those Interested In PyTorch