[Python知识库]pytorch gpu 加速

在GPU上运行代码

判断GPU是否可用

torch.cuda.is_available()

定义一个device对象

device=torch.device('cuda' if torch.cuda.is_available() else 'gpu')
print('device:',device)

用法展示

a=torch.Tensor([1,2,3])
print(a)
print(a.device)
a.to(device)
print(a.device)
a=a.to(device)
print(a)



# 推荐写法
b=torch.Tensor([1,2,3]).to(device)
print('b=',b)

tensor([ 1.,  2.,  3.])
cpu
cpu
tensor([ 1.,  2.,  3.], device='cuda:0')
b= tensor([ 1.,  2.,  3.], device='cuda:0')

模型to gpu 参数在自动转到 gpu

gpu上的对象，运算得到的结果也是gpu上的。

注意：

1、自定义得参数和数据，需要转化为cuda支持得tensor

2、model 需要转化为cuda 支持得tensor

3、执行结果需要和cpu的tensor计算的时候

? a.cup()把cuda 的tensor 转化成 cpu 的tensor

a.cpu()
print(a)
a=a.cpu()
print(a.device)

tensor([ 1.,  2.,  3.], device='cuda:0')
cpu