CUDA semantics¶
torch.cuda
keeps track of currently selected GPU, and all CUDA tensors
you allocate will be created on it. The selected device can be changed with a
torch.cuda.device
context manager.
However, once a tensor is allocated, you can do operations on it irrespectively of your selected device, and the results will be always placed in on the same device as the tensor.
Cross-GPU operations are not allowed by default, with the only exception of
copy_()
. Unless you enable peer-to-peer memory accesses
any attempts to launch ops on tensors spread across different devices will
raise an error.
Below you can find a small example showcasing this:
x = torch.cuda.FloatTensor(1)
# x.get_device() == 0
y = torch.FloatTensor(1).cuda()
# y.get_device() == 0
with torch.cuda.device(1):
# allocates a tensor on GPU 1
a = torch.cuda.FloatTensor(1)
# transfers a tensor from CPU to GPU 1
b = torch.FloatTensor(1).cuda()
# a.get_device() == b.get_device() == 1
c = a + b
# c.get_device() == 1
z = x + y
# z.get_device() == 0
# even within a context, you can give a GPU id to the .cuda call
d = torch.randn(2).cuda(2)
# d.get_device() == 2
Best practices¶
Use pinned memory buffers¶
Host to GPU copies are much faster when they originate from pinned (page-locked)
memory. CPU tensors and storages expose a pin_memory()
method, that returns a copy of the object, with data put in a pinned region.
Also, once you pin a tensor or storage, you can use asynchronous GPU copies.
Just pass an additional async=True
argument to a cuda()
call. This can be used to overlap data transfers with computation.
You can make the DataLoader
return batches placed in
pinned memory by passing pin_memory=True
to its constructor.
Use nn.DataParallel instead of multiprocessing¶
Most use cases involving batched input and multiple GPUs should default to using
DataParallel
to utilize more than one GPU. Even with the GIL,
a single python process can saturate multiple GPUs.
As of version 0.1.9, large numbers of GPUs (8+) might not be fully utilized. However, this is a known issue that is under active development. As always, test your use case.
There are significant caveats to using CUDA models with
multiprocessing
; unless care is taken to meet the data handling
requirements exactly, it is likely that your program will have incorrect or
undefined behavior.