View the graphics card information
import torch
from torch import nn
from d2l import torch as d2l
Computing Devices
def cpu():
"""Get the CPU device."""
return torch.device('cpu')
def gpu(i=0):
"""Get a GPU device."""
return torch.device(f'cuda:{i}')
cpu(), gpu(), gpu(1)
(device(type='cpu'), device(type='cuda', index=0), device(type='cuda', index=1))
Query the number of available GPUs
def num_gpus():
"""Get the number of available GPUs."""
return torch.cuda.device_count()
num_gpus()
2
Define two convenient functions that allow us to run code even if the requested GPUs do not exist
def try_gpu(i=0):
"""Return gpu(i) if exists, otherwise return cpu()."""
if num_gpus() >= i + 1:
return gpu(i)
return cpu()
def try_all_gpus():
"""Return all available GPUs, or [cpu(),] if no GPU exists."""
return [gpu(i) for i in range(num_gpus())]
try_gpu(), try_gpu(10), try_all_gpus()
(device(type='cuda', index=0), device(type='cpu'), [device(type='cuda', index=0), device(type='cuda', index=1)])
Query the device where the tensor is located
x = torch.tensor([1, 2, 3])
x.device
device(type='cpu')
Store a tensor on the GPU
X = torch.ones(2, 3, device=try_gpu())
X
tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:0')
Create a random tensor, Y
, on the second GPU
Y = torch.rand(2, 3, device=try_gpu(1))
Y
tensor([[0.0022, 0.5723, 0.2890], [0.1456, 0.3537, 0.7359]], device='cuda:1')
If we want to compute X + Y
,
we need to decide where to perform this operation
Z = X.cuda(1)
print(X)
print(Z)
tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:0') tensor([[1., 1., 1.], [1., 1., 1.]], device='cuda:1')
The data (both Z
and Y
) are on the same GPU), we can add them up
Y + Z
tensor([[1.0022, 1.5723, 1.2890], [1.1456, 1.3537, 1.7359]], device='cuda:1')
Z.cuda(1) is Z
True
Neural Networks and GPUs
net = nn.Sequential(nn.LazyLinear(1))
net = net.to(device=try_gpu())
net(X)
tensor([[0.7802], [0.7802]], device='cuda:0', grad_fn=<AddmmBackward0>)
Confirm that the model parameters are stored on the same GPU
net[0].weight.data.device
device(type='cuda', index=0)
@d2l.add_to_class(d2l.Trainer)
def __init__(self, max_epochs, num_gpus=0, gradient_clip_val=0):
self.save_hyperparameters()
self.gpus = [d2l.gpu(i) for i in range(min(num_gpus, d2l.num_gpus()))]
@d2l.add_to_class(d2l.Trainer)
def prepare_batch(self, batch):
if self.gpus:
batch = [a.to(self.gpus[0]) for a in batch]
return batch
@d2l.add_to_class(d2l.Trainer)
def prepare_model(self, model):
model.trainer = self
model.board.xlim = [0, self.max_epochs]
if self.gpus:
model.to(self.gpus[0])
self.model = model