GPU Basics

By the end of this lesson you will understand why deep learning runs so much faster on a GPU, how to write PyTorch code that uses one when available and falls back to the CPU when not, and how to deal with the two GPU problems everyone meets: device mismatches and running out of memory.

You do not need a GPU to finish this course. Every example works on a CPU, just more slowly for big models. But real projects almost always train on GPUs, so you should write "device-ready" code from the start.

Why GPUs are fast for deep learning

A CPU (central processing unit) has a handful of very capable cores. It is great at complicated tasks with many decisions. A GPU (graphics processing unit) has thousands of simpler cores. It was built to colour millions of screen pixels at once, and it is great at doing the same simple calculation on huge amounts of data.

Think of a gurudwara langar kitchen. One expert chef could cook any dish, but to make 10,000 rotis you want hundreds of volunteers each rolling and cooking at the same time. Neural network training is mostly rotis: enormous matrix multiplications where every output number is an independent sum of products.

Step 1: count the work in one training step. You can estimate how many multiply-and-add operations a network needs.

def dense_flops(sizes, batch):
    """Multiply-adds for one forward pass of a stack of dense layers."""
    return sum(batch * n_in * n_out for n_in, n_out in zip(sizes[:-1], sizes[1:]))

small = dense_flops([784, 256, 10], batch=64)
big = dense_flops([3 * 224 * 224, 4096, 4096, 1000], batch=64)
print(f"small digit MLP: {small:,} multiply-adds per batch")
print(f"big image MLP:   {big:,} multiply-adds per batch")

# training is roughly 3x the forward pass (forward + backward for inputs and weights)
for name, work in [("small", small), ("big", big)]:
    print(f"{name}: about {3 * work / 1e6:,.0f} million operations per training step")

small digit MLP: 13,008,896 multiply-adds per batch
big image MLP:   40,795,897,856 multiply-adds per batch
small: about 39 million operations per training step
big: about 122,388 million operations per training step

The small network needs about 39 million operations per training step, which any laptop handles easily. The big one needs over 120,000 million (more than 120 billion) for every batch. Each one is simple and independent, which is exactly what thousands of GPU cores do well. That is why a GPU can train such a model many times faster than a CPU.

The kinds of accelerators

PyTorch can use several kinds of hardware. Each has its own device name:

CUDA is NVIDIA's platform for running general computations on its GPUs. The standard pip install torch command on Windows installs a CPU-only build. The pytorch.org "Get Started" page gives the exact command for a CUDA build that matches your system.

Step 2: choose the device once

Write one line at the top of your script that picks the best device, then use that variable everywhere. Recent PyTorch versions include torch.accelerator, which detects CUDA, MPS, XPU and others in one call:

import torch

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print("using device:", device)
print("CUDA available?", torch.cuda.is_available())

using device: cpu
CUDA available? False

This output came from a laptop with no supported GPU, so the code fell back to the CPU without any change. On a machine with an NVIDIA GPU it would print cuda. In older code you will see the CUDA-only version: device = "cuda" if torch.cuda.is_available() else "cpu".

Step 3: move the model and the data

A tensor lives on one device. A calculation can only combine tensors on the same device. So you move:

1. the model, once, right after creating it: model = model.to(device);
2. each batch, inside the loop: images, labels = images.to(device), labels.to(device).

import torch
from torch import nn

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
model = nn.Linear(4, 2).to(device)                 # moves every parameter
x = torch.randn(3, 4).to(device)                   # moves this batch
print("weight on:", model.weight.device, "| input on:", x.device)
print("output on:", model(x).device)

result = model(x).detach().cpu().numpy()           # back to NumPy: detach, then cpu
print("as NumPy:", result.shape)

weight on: cpu | input on: cpu
output on: cpu
as NumPy: (3, 2)

If you forget to move the batch, a GPU run fails with an error saying it expected all tensors to be on the same device but found two devices, such as cuda:0 and cpu. The fix is always a missing .to(device). Note too that .to() on a tensor returns a new tensor, so you must assign the result: x.to(device) on its own line does nothing useful. For a model, .to() moves it in place, but writing model = model.to(device) is a safe habit.

To get results back for NumPy, matplotlib or printing, use .detach().cpu(), as in the last line.

GPU memory and "out of memory"

A GPU has its own memory, separate from your computer's RAM, often between 4 and 24 GB on consumer cards. During training it holds the model's weights, their gradients, the optimiser's state (Adam keeps two extra numbers per weight) and the activations saved for the backward pass. Activations grow with the batch size.

When it fills up, PyTorch raises a "CUDA out of memory" error. Fixes, in the order to try them:

1. Lower the batch size. Halving it roughly halves the activation memory.
2. Use mixed precision (next section), which stores many values in 16 bits instead of 32.
3. Make sure evaluation runs under torch.no_grad(), so no activations are kept.
4. Do not keep tensors that hold the graph, such as appending loss to a list. Append loss.item() instead.
5. Use a smaller model or a smaller input size.

Mixed precision

By default, weights and activations are 32-bit floats (float32). Modern GPUs compute much faster with 16-bit floats, either float16 or bfloat16. Mixed precision runs most operations in 16 bits while keeping sensitive ones, like the loss and the weight updates, in 32 bits. It often speeds training up and cuts memory use, usually with no loss in accuracy.

In PyTorch you wrap the forward pass in torch.autocast. With float16 you also use a GradScaler, which multiplies the loss before backward() so tiny gradients do not round down to zero in 16 bits. bfloat16 has the same range as float32, so it does not need a scaler. A typical CUDA training step looks like this (plain code, it needs an NVIDIA GPU):

scaler = torch.amp.GradScaler("cuda")

for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(images), labels)      # forward pass in mixed precision
    optimizer.zero_grad()
    scaler.scale(loss).backward()                  # scaled backward pass
    scaler.step(optimizer)                         # unscale, then update
    scaler.update()

Autocast also works on CPUs with bfloat16. Here it is on this laptop's CPU:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Linear(8, 4)
x = torch.randn(2, 8)
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    out = model(x)
print("inside autocast:", out.dtype)
print("weights stay:   ", model.weight.dtype)

inside autocast: torch.bfloat16
weights stay:    torch.float32

The layer's output was computed in bfloat16, but the stored weights stayed float32. That is the "mixed" in mixed precision.

Free GPUs for learning

If you have no NVIDIA GPU, cloud notebooks such as Google Colab and Kaggle Notebooks offer limited free GPU time, which is enough for every project in this course. Their limits change over time, so check their current terms. On your own machine with an NVIDIA card, the nvidia-smi command in a terminal shows the GPU, its memory and what is using it.

Common mistakes

• Writing x.to(device) without assigning it. Tensors are not moved in place; write x = x.to(device).
• Moving the model but not the batch (or the other way round). Every input to an operation must be on the same device.
• Creating the optimiser before model.to(device). Create it after moving the model.
• Calling .numpy() on a GPU tensor. Use .detach().cpu().numpy().
• Collecting loss tensors in a list. This keeps every graph in memory. Store loss.item().

Key points

• GPUs have thousands of simple cores that do huge matrix multiplications in parallel.
• Pick the device once (torch.accelerator or torch.cuda.is_available()) and fall back to "cpu".
• Move the model once and every batch inside the loop with .to(device); assign the result for tensors.
• Out-of-memory errors are fixed first by smaller batches, mixed precision and torch.no_grad() for evaluation.
• Mixed precision uses torch.autocast, plus GradScaler when training in float16.

import torch
from torch import nn

torch.manual_seed(0)
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)    # after .to(device)

x = torch.randn(16, 10).to(device)
y = torch.randn(16, 1).to(device)
loss = nn.functional.mse_loss(model(x), y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"loss: {loss.item():.4f}")
print("first weight on:", model[0].weight.device)

loss: 1.5164
first weight on: cpu