在进行深度学习训练模型时,对于计算量小一些的模型,是可以在CPU上进行的。但是当计算量比较大时,我们希望利用GPU并行计算的能力去加快训练的速度。
查看GPU版本号
import torch
print(torch.__version__) # 查看torch当前版本号
print(torch.version.cuda) # 编译当前版本的torch使用的cuda版本号
print(torch.cuda.is_available()) # 查看当前cuda是否可用于当前版本的Torch,如果输出True,则表示可用
查看GPU数量
def try_gpu(i=0):
"""如果存在,则返回gpu(i),否则返回cpu()"""
if torch.cuda.device_count() >= i + 1:
return torch.device(f'cuda:{i}')
return torch.device('cpu')
def try_all_gpus():
"""返回所有可用的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')]
# 0号GPU是否存在,10号GPU是否存在
try_gpu(), try_gpu(10), try_all_gpus()
指定GPU
import torch
from torch import nn
torch.device('gpu'), torch.cuda.device('cuda'), torch.cuda.device('cuda:1')
GPU计算张量
# 创建一个张量Y在1号GPU
Y = torch.rand(2, 3, device=try_gpu(1))
Z = X.cuda(1) # 将X的内容复制在1号GPU的Z
print(X)
print(Z)
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:0')
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:1')
指定GPU计算神经网络模型
from torch import nn
net = nn.Linear(3, 1)
print(list(net.parameters())[0].device) # cpu
net.cuda(0)
print(list(net.parameters())[0].device) # cuda:0
net2 = nn.Linear(3, 1, device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
print(list(net2.parameters())[0].device) # cuda:0
net3 = nn.Linear(3, 1).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
print(list(net3.parameters())[0].device) # cuda:0
net3 = net3.cpu()
print(list(net3.parameters())[0].device) # cpu