Using torch to compare CPU/GPU speeds - imageprocessing-awesome

stough 202-

Preview: Using [torch](https://pytorch.org/docs/stable/cuda.html) to compare CPU/GPU speeds

The Graphics Processing Unit is a common coprocessor designed to do parallel floating point arithmetic. In the past this was computer graphics, but this massively parallel math is useful in all scientific computation.

Also, going to use jupyterlab magic commands.

The neural networks in this chapter do enormous numbers of multiply-adds, which is exactly the kind of work GPUs are built for. Here we time one big batch of matrix multiplications three ways, in numpy, in torch on the CPU, and in torch on the GPU, to see what the fuss is about, and to pick up a couple of gotchas about measuring GPU speed. If your machine has no CUDA GPU (or you’re in Colab without a GPU runtime), the “GPU” cells will simply run on the CPU again.

import torch
import torch.cuda as cuda
import numpy as np

import time
from IPython.display import display, Markdown

cuda.is_available()

# Use the GPU if there is one; otherwise everything below still runs, on the CPU.
device = torch.device('cuda' if cuda.is_available() else 'cpu')
print('Using', device)
if cuda.is_available():
    print(cuda.get_device_name(0))

if cuda.is_available():
    props = cuda.get_device_properties(0)
    print(props)
    print(f'{props.total_memory / 1024**3:.1f} GB of GPU memory, {props.multi_processor_count} streaming multiprocessors')

We’ll do a large matrix multiply operation¶

in numpy, torch, and torch on the GPU.

A holds BATCH matrices of size $1000 \times 200$ , and B holds BATCH matrices of size $200 \times 1000$ . np.matmul multiplies them pairwise, giving BATCH matrices of size $1000 \times 1000$ . Each product takes $1000 \cdot 1000 \cdot 200 = 2 \times 10^8$ multiply-adds. Keep an eye on memory, though: in 8-byte float64, C alone is BATCH $\times$ 8 MB. The GPU has to hold A, B and C all at once, so lower BATCH if you run out of memory.

BATCH = 20    # number of matrix pairs; C will take BATCH * 8 MB

A = np.random.rand(BATCH,1000,200)
B = np.random.rand(BATCH,200,1000)

(8*A.size)/(1024**2)

C = np.matmul(A,B)
C.shape

8*C.size/(1024**2)

%%timeit -n 5 -r 4
# C = np.matmul(A,B)
np.matmul(A,B, out=C)

%%timeit -n 5 -r 4 runs the cell 5 times per trial for 4 trials and reports the average. Using out=C writes into the already-allocated C, so we time the arithmetic rather than allocating a new C each time.

Test in Torch¶

At_cpu = torch.from_numpy(A)
Bt_cpu = torch.from_numpy(B)

At_cpu.is_cuda

torch.from_numpy doesn’t copy anything: the tensor shares memory with the numpy array. numpy and torch on the CPU both hand the work to optimized, multi-core linear algebra code, but how well each handles a batch of matrices differs. On the machine this was written on, torch’s CPU version was about 8 times faster than numpy’s for this batched multiply. So even “just use the CPU” depends on the library.

Ct_cpu = torch.matmul(At_cpu, Bt_cpu)
Ct_cpu.shape

%%timeit -n 5 -r 4
# Ct_cpu = torch.matmul(At_cpu, Bt_cpu)
torch.matmul(At_cpu, Bt_cpu, out = Ct_cpu)

Now test in torch, on the GPU¶

First, the data has to be copied to the GPU’s own memory. That copy across the PCIe bus is not free, so let’s time it too.

st = time.time()
At_gpu = At_cpu.to(device)
Bt_gpu = Bt_cpu.to(device)
if cuda.is_available():
    cuda.synchronize()
print(f'copying A and B to {device} took {time.time() - st:.3f} s')

Ct_gpu = torch.zeros_like(Ct_cpu).to(device)

At_gpu.is_cuda

%%timeit -n 2 -r 3
torch.matmul(At_gpu, Bt_gpu, out=Ct_gpu)
if cuda.is_available():
    cuda.synchronize()   # wait for the GPU to actually finish

Note the cuda.synchronize(). GPU operations in torch are asynchronous: torch.matmul on a GPU tensor just queues the work and returns immediately, while the GPU crunches in the background. Without waiting for it to finish, we would mostly be timing how fast Python can queue up work, which can make the GPU look impossibly fast.

Single precision¶

How much faster the GPU is depends heavily on the number type. Most consumer and workstation GPUs are built for 32-bit floats (what graphics and deep learning use) and run 64-bit float64 math at a small fraction of that speed, as little as $\frac{1}{32}$ or $\frac{1}{64}$ . Let’s repeat the comparison in float32, which is what neural networks normally use.

A32, B32 = At_cpu.float(), Bt_cpu.float()
C32 = torch.matmul(A32, B32)
A32_gpu, B32_gpu, C32_gpu = A32.to(device), B32.to(device), C32.to(device)

def best_time(f, reps=5):
    times = []
    for _ in range(reps):
        st = time.time()
        f()
        if cuda.is_available():
            cuda.synchronize()
        times.append(time.time() - st)
    return min(times)

t64_cpu = best_time(lambda: torch.matmul(At_cpu, Bt_cpu, out=Ct_cpu))
t64_gpu = best_time(lambda: torch.matmul(At_gpu, Bt_gpu, out=Ct_gpu), reps=2)
t32_cpu = best_time(lambda: torch.matmul(A32, B32, out=C32))
t32_gpu = best_time(lambda: torch.matmul(A32_gpu, B32_gpu, out=C32_gpu))

print(f'float64: cpu {1000*t64_cpu:8.1f} ms, {device} {1000*t64_gpu:8.1f} ms, speedup {t64_cpu/t64_gpu:5.1f}x')
print(f'float32: cpu {1000*t32_cpu:8.1f} ms, {device} {1000*t32_gpu:8.1f} ms, speedup {t32_cpu/t32_gpu:5.1f}x')

Try again, with our own timing.¶

math expressions in markdown

times records the cumulative time after each iteration, so it should grow roughly linearly. Each iteration also zeroes Ct_gpu first, which is cheap next to the multiply.

NUMITERS = 10
times = []

st = time.time()

for i in range(NUMITERS):
    torch.zeros(Ct_gpu.shape, out=Ct_gpu)
    torch.matmul(At_gpu, Bt_gpu, out=Ct_gpu)
    if cuda.is_available():
        cuda.synchronize()
    times.append(time.time() - st)

et = time.time()

# Why not be more complicated...
display(Markdown(rf'{NUMITERS} iters took {1000*(et-st)/NUMITERS:.2f} ms per iteration.'))

times

Takeaways¶

How much a GPU helps depends enormously on the hardware and the number type. Data-center and desktop GPUs can beat a CPU by 10 to 100 times on float32 matrix math. A small laptop GPU may only roughly match a good multi-core CPU, and in float64 it can be dramatically slower. On the laptop this notebook was written on, float64 ran about 10 times slower on its GPU than on the CPU, while float32 ran about 3 times faster. Deep learning uses float32 (or even 16-bit) numbers for exactly this reason.
Moving data between CPU and GPU memory costs real time. Copy data to the GPU once and keep it there for many operations, which is what the training loops in the following notebooks do with the network’s weights.
Always synchronize before reading a clock when timing GPU code.

Run this notebook on your own machine, and in a Colab GPU runtime, and compare.