working-with-gpus

Duration: 8 min

This module delves into the essentials of utilizing GPUs for deep learning with PyTorch. Understanding how to leverage GPUs can significantly speed up model training, making it crucial for efficient and scalable deep learning projects.

Checking GPU Availability

Before utilizing a GPU, it's important to check if one is available. PyTorch provides a straightforward way to check for GPU availability using the torch.cuda.is_available() function. This function returns a boolean indicating whether CUDA is available.

import torch

# Check if CUDA (GPU support) is available
is_gpu_available = torch.cuda.is_available()

print(f'Is CUDA available: {is_gpu_available}')

# If CUDA is available, print the number of GPUs
if is_gpu_available:
    print(f'Number of GPUs: {torch.cuda.device_count()}')

Try it in Google Colab:

Is CUDA available: True
Number of GPUs: 1

Moving Data to GPU

Once you've confirmed that a GPU is available, the next step is to move your data and models to the GPU. PyTorch allows you to easily transfer tensors and models to the GPU using the .to(device) method, where 'device' is a torch.device object specifying the GPU.

import torch

# Check if CUDA is available
is_gpu_available = torch.cuda.is_available()
device = torch.device('cuda' if is_gpu_available else 'cpu')

# Create a tensor and move it to the GPU
tensor_cpu = torch.tensor([1.0, 2.0, 3.0])
tensor_gpu = tensor_cpu.to(device)

print(f'Tensor on {device}:\n{tensor_gpu}')

💡 Tip: Always ensure that both your model and your data are on the same device to avoid errors during computation.