GPU相关
查看计算机有多少台GPU
import torch
from torch import nn
torch.cuda.device_count()
输出
1
定义一些常用函数
def try_gpu(i=0): #@save
"""如果存在,则返回gpu(i),否则返回cpu()"""
if torch.cuda.device_count() >= i + 1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
def try_all_gpus(): #@save
"""返回所有可用的GPU,如果没有GPU,则返回[cpu(),]"""
devices = [
torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())
]
return devices if devices else [torch.device('cpu')]
try_gpu(),try_gpu(10),try_all_gpus()
try_gpu()有就返回GPU,没有就返回CPU
try_all_gpus()返回所有GPU
输出
(device(type='cuda', index=0),
device(type='cpu'),
[device(type='cuda', index=0)])
在GPU上使用张量
x = torch.tensor([1, 2, 3])
x.device
输出
device(type='cpu')
默认是在CPU上运行的
指定在GPU上运行
X = torch.ones(2, 3, device = try_gpu())
X
输出
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:0')
张量的计算
x + X
RuntimeError Traceback (most recent call last)
<ipython-input-20-cade78b2014d> in <module>
----> 1 x + X
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
可以看到不再一起是不能做运算的
Y = torch.rand(2, 3, device=try_gpu())
Y
输出
tensor([[0.8074, 0.9579, 0.0942],
[0.1074, 0.5444, 0.1546]], device='cuda:0')
X + Y
输出
tensor([[1.8074, 1.9579, 1.0942],
[1.1074, 1.5444, 1.1546]], device='cuda:0')
在GPU上使用神经网络
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device = try_gpu())
net(X)
输出
tensor([[0.0069],
[0.0069]], device='cuda:0', grad_fn=<AddmmBackward0>)
查看参数位置
net[0].weight.data.device
输出
device(type='cuda', index=0)