23 KiB
23 KiB
In [1]:
import torch
from torch import nn
from d2l import torch as d2lIn [2]:
def cpu(): #@save
"""Get the CPU device."""
return torch.device('cpu')
def gpu(i=0): #@save
"""Get a GPU device."""
return torch.device(f'cuda:{i}')
cpu(), gpu(), gpu(1)Out [2]:
(device(type='cpu'), device(type='cuda', index=0), device(type='cuda', index=1))
In [3]:
def num_gpus(): #@save
"""Get the number of available GPUs."""
return torch.cuda.device_count()
num_gpus()Out [3]:
2
In [4]:
def try_gpu(i=0): #@save
"""Return gpu(i) if exists, otherwise return cpu()."""
if num_gpus() >= i + 1:
return gpu(i)
return cpu()
def try_all_gpus(): #@save
"""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()Out [4]:
(device(type='cuda', index=0), device(type='cpu'), [device(type='cuda', index=0), device(type='cuda', index=1)])
In [5]:
x = torch.tensor([1, 2, 3])
x.deviceOut [5]:
device(type='cpu')
In [6]:
X = torch.ones(2, 3, device=try_gpu())
XOut [6]:
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:0')In [7]:
Y = torch.rand(2, 3, device=try_gpu(1))
YOut [7]:
tensor([[0.0022, 0.5723, 0.2890],
[0.1456, 0.3537, 0.7359]], device='cuda:1')In [8]:
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')
In [9]:
Y + ZOut [9]:
tensor([[1.0022, 1.5723, 1.2890],
[1.1456, 1.3537, 1.7359]], device='cuda:1')In [10]:
Z.cuda(1) is ZOut [10]:
True
In [11]:
net = nn.Sequential(nn.LazyLinear(1))
net = net.to(device=try_gpu())In [12]:
net(X)Out [12]:
tensor([[0.7802],
[0.7802]], device='cuda:0', grad_fn=<AddmmBackward0>)In [13]:
net[0].weight.data.deviceOut [13]:
device(type='cuda', index=0)
In [14]:
@d2l.add_to_class(d2l.Trainer) #@save
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) #@save
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) #@save
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