Concept units

PyTorch tensors on the GPU

Intermediate

Read these first if they are new to you

The idea in one paragraph

A PyTorch tensor is an array, like a NumPy array, with one more property: the device it lives on. On the CPU, the computer's main processor, it works much like NumPy. On a GPU, a graphics card, thousands of simple cores apply the same step to thousands of numbers at once. The catch is that the GPU has its own memory. Data must be copied there before the work and copied back after it, and that copy is slow compared with the work. So programs move data to the GPU once and do many steps there.

A tensor is an array that knows which device holds it

PyTorch is a Python library for machine learning. Its array type is the tensor. A tensor has a shape and a dtype, just like a NumPy array, and most NumPy operations have a PyTorch twin. It also has a device: cpu, or cuda for an NVIDIA GPU. CUDA is NVIDIA's system for running general programs on its graphics cards.

The method .to(device) returns a copy of the tensor on that device. .cpu() copies it back, and .numpy() turns a CPU tensor into a NumPy array. Before you read the output below, predict the two dtypes on the first line.

tensor.py

import numpy as np
import torch

a = np.array([[0.0, 0.0, 0.0], [2.0, 4.0, 0.0]])
x = torch.from_numpy(a)                 # from a NumPy array
y = torch.tensor([[1.0, 2.0, 3.0]])     # from a Python list
print(x.dtype, y.dtype)
print(x.shape, x.device)

# Pick "cuda" when PyTorch finds a GPU, and "cpu" otherwise.
device = "cuda" if torch.cuda.is_available() else "cpu"
x = x.float().to(device)                # a float32 copy on the GPU, if there is one
z = x * 10                              # runs where x lives
print(z.device)
print(z.cpu().numpy())                  # copy back to the CPU, then to NumPy

torch.float64 torch.float32
torch.Size([2, 3]) cpu
cpu
[[ 0.  0.  0.]
 [20. 40.  0.]]

NumPy makes 64-bit decimals by default, and torch.from_numpy keeps that type. PyTorch makes 32-bit decimals, float32, by default. Neural networks almost always use float32 or smaller types, because GPUs handle them much faster. .float() converts to float32.

The second line of output prints the device before the move, so it says cpu on any computer. The output above came from a computer without a GPU, so z stayed on the CPU too, and the third line says cpu. On a computer with a GPU, the third line says cuda:0, the first GPU. Every other line stays the same.

A GPU applies one step to thousands of numbers at once

A CPU has a few powerful cores, the parts that carry out instructions. A desktop CPU typically has from 4 to 24. Each core can follow its own complicated program. A GPU has thousands of simpler cores that work in groups, all doing the same step on different numbers. Adding 1 to a million numbers is exactly that kind of job, and so are most steps of a neural network.

Two processors, two memories, and a narrow link between them
CPU a few strong cores main memory (RAM) GPU thousands of simple cores GPU memory copy: slow PCIe bus
Each processor reads its own memory fast. Moving data across the link between them is much slower, so a tensor stays on one device while it is worked on.

Every tensor in one operation must be on the same device. If x is on the GPU and y is on the CPU, x + y stops with an error that says all tensors were expected to be on the same device. The fix is to move one of them with .to(...).

The copy can take longer than the work

Worked example: is one step worth a trip to the GPU?

Take a million 3D points as float32, and add 1 to each number.

  1. Size: 1,000,000 points × 3 numbers × 4 bytes = 12,000,000 bytes, or 12 MB.
  2. Copy to the GPU: a PCIe 4.0 link moves roughly 25 GB per second in practice. 12 MB ÷ 25 GB/s ≈ 0.5 ms.
  3. Work on the GPU: the step reads 12 MB and writes 12 MB. An NVIDIA A100 GPU's memory is rated at 1,555 GB per second or more, and newer GPUs are faster. 24 MB ÷ 1,500 GB/s ≈ 0.016 ms.
  4. Copy back: another 0.5 ms.
  5. Work on the CPU instead, with no copies: on the computer that ran this page's snippets, one CPU core did the same step in about 0.4 ms.

The GPU route costs about 1 ms, almost all of it copying. The CPU does the whole step in less than half that time. For one step, the GPU is not worth it. For a neural network that does thousands of steps on the same data, it is.

size.py

import torch

clouds = torch.zeros(8, 2048, 3)        # all zeros: 8 point clouds, 2048 points each
print(clouds.shape, clouds.dtype)
print(clouds.nelement(), "numbers of", clouds.element_size(), "bytes each")
print(clouds.nelement() * clouds.element_size() / 1_000_000, "MB")

torch.Size([8, 2048, 3]) torch.float32
49152 numbers of 4 bytes each
0.196608 MB
A race: 256 numbers, several steps each

The CPU here does 8 numbers per tick. The GPU does all 256 per tick, but each copy takes 40 ticks. At 1 step, which finishes first? Predict, press Run, then drag the slider until the GPU wins.

Press Run.
Made-up speeds, real shape. Total ticks are CPU = 32 × steps and GPU = 40 + steps + 40. The GPU wins once there are enough steps to pay for the two copies.
These numbers are rough

Link and memory speeds vary several times between machines, and the race uses invented speeds. A PCIe 4.0 link is rated at about 32 GB per second each way and reaches about 25 GB per second in practice. A PCIe 5.0 link doubles both. Copies from ordinary memory, rather than memory set aside for copying, are often slower. CPUs also do some work in parallel: each core can apply one instruction to several numbers. The shape of the trade-off holds: copying is slow, and a GPU pays off when it does many steps per copy.

GPU work runs in the background

PyTorch hands GPU work to a queue and returns at once, before the work is done. A Python timer around a GPU step therefore measures almost nothing. Call torch.cuda.synchronize(), which waits for the queue to empty, before you read the clock. Copying a result back with .cpu() or .item() also waits.

Check yourself

  1. A model's weights are on the GPU. You call it on a tensor batch made with torch.from_numpy. It fails with a device error. What one change fixes it, and what second problem might remain?

    Answer

    Move the input: batch = batch.to("cuda"). The tensor may also still be float64, while the weights are float32. Use batch.float().to("cuda").

  2. How many megabytes does a float32 tensor of shape (16, 4096, 3) take?

    Answer

    16 × 4096 × 3 = 196,608 numbers, times 4 bytes = 786,432 bytes, or about 0.79 MB.

  3. In the race, with 3 steps, how many ticks does each side take, and who wins?

    Answer

    CPU: 32 × 3 = 96 ticks. GPU: 40 + 3 + 40 = 83 ticks. The GPU wins. At 2 steps the CPU wins, 64 to 82.

Go further

Optional background.

Where this shows up

Units that build on this one