Skip to main content
W&B Fully Connected
>
Articles

How To Use GPU with PyTorch

A short tutorial on using GPUs for your deep learning models with PyTorch, from checking availability to visualizing usable.
Ayush Thakur
Created on November 20|Last edited on January 3

Sections

Introduction

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.
Before continuing and if you haven't already, you may want to check if Pytorch is using your GPU.

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_cuda
    False
  • 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_cuda
    True
  • 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.



Run set
1


Summary

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.

Try Weights & Biases

Weights & Biases helps you keep track of your machine learning experiments. Try 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 or run 2 quick experiments on Replit and see how W&B can help organise your work foloow the instructions below:
Instructions:
  1. Click the green "Run" button below (the first time you click Run, Replit will take approx 30-45 seconds to allocate a machine)
  2. Follow the prompts in the terminal window (the bottom right pane below)
  3. You can resize the terminal window (bottom right) for a larger view




Sadegh Pouriyan
Sadegh Pouriyan •  
You imported torch and checked whether Cuda is available or not. Which version of PyTorch were you installed? GPU version or cpuonly version?
1 reply
Tags: Beginner, Domain Agnostic, PyTorch, Tutorial, Plots, Chum here

Created with ❤️ on Weights & Biases.

https://wandb.ai/wandb/common-ml-errors/reports/How-To-Use-GPU-with-PyTorch---VmlldzozMzAxMDk
Made with Weights & Biases. Sign up or log in to create reports like this one.
Iterate on AI agents and models faster. Try Weights & Biases today.
Sign upTry W&B now