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
Ayush Thakur



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

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.

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.
Get started in 5 minutes.

Recommended Reading For Those Interested In PyTorch

Report Gallery