12 KiB
12 KiB
In [1]:
import torch
from d2l import torch as d2lIn [2]:
devices = d2l.try_all_gpus()
def run(x):
return [x.mm(x) for _ in range(50)]
x_gpu1 = torch.rand(size=(4000, 4000), device=devices[0])
x_gpu2 = torch.rand(size=(4000, 4000), device=devices[1])In [3]:
run(x_gpu1)
run(x_gpu2) # Warm-up all devices
torch.cuda.synchronize(devices[0])
torch.cuda.synchronize(devices[1])
with d2l.Benchmark('GPU1 time'):
run(x_gpu1)
torch.cuda.synchronize(devices[0])
with d2l.Benchmark('GPU2 time'):
run(x_gpu2)
torch.cuda.synchronize(devices[1])GPU1 time: 0.4660 sec
GPU2 time: 0.4510 sec
In [4]:
with d2l.Benchmark('GPU1 & GPU2'):
run(x_gpu1)
run(x_gpu2)
torch.cuda.synchronize()GPU1 & GPU2: 0.4659 sec
In [5]:
def copy_to_cpu(x, non_blocking=False):
return [y.to('cpu', non_blocking=non_blocking) for y in x]
with d2l.Benchmark('Run on GPU1'):
y = run(x_gpu1)
torch.cuda.synchronize()
with d2l.Benchmark('Copy to CPU'):
y_cpu = copy_to_cpu(y)
torch.cuda.synchronize()Run on GPU1: 0.4656 sec
Copy to CPU: 2.3125 sec
In [6]:
with d2l.Benchmark('Run on GPU1 and copy to CPU'):
y = run(x_gpu1)
y_cpu = copy_to_cpu(y, True)
torch.cuda.synchronize()Run on GPU1 and copy to CPU: 1.6907 sec