当前位置:   article > 正文

pytorch采用GPU加速方法_pytorch gpu加速

pytorch gpu加速

在进行深度学习训练模型时,对于计算量小一些的模型,是可以在CPU上进行的。但是当计算量比较大时,我们希望利用GPU并行计算的能力去加快训练的速度。

查看GPU版本号

  1. import torch
  2. print(torch.__version__)  # 查看torch当前版本号
  3. print(torch.version.cuda)  # 编译当前版本的torch使用的cuda版本号
  4. print(torch.cuda.is_available())  # 查看当前cuda是否可用于当前版本的Torch,如果输出True,则表示可用

查看GPU数量

  1. def try_gpu(i=0):
  2. """如果存在,则返回gpu(i),否则返回cpu()"""
  3. if torch.cuda.device_count() >= i + 1:
  4. return torch.device(f'cuda:{i}')
  5. return torch.device('cpu')
  6. def try_all_gpus():
  7. """返回所有可用的GPU,如果没有GPU,则返回[cpu(),]"""
  8. devices = [torch.device(f'cuda:{i}')
  9. for i in range(torch.cuda.device_count())]
  10. return devices if devices else [torch.device('cpu')]
  11. # 0号GPU是否存在,10号GPU是否存在
  12. try_gpu(), try_gpu(10), try_all_gpus()


指定GPU

  1. import torch
  2. from torch import nn
  3. torch.device('gpu'), torch.cuda.device('cuda'), torch.cuda.device('cuda:1')

GPU计算张量

  1. # 创建一个张量Y在1号GPU
  2. Y = torch.rand(2, 3, device=try_gpu(1))
  3. Z = X.cuda(1) # 将X的内容复制在1号GPU的Z
  4. print(X)
  5. print(Z)
  6. tensor([[1., 1., 1.],
  7. [1., 1., 1.]], device='cuda:0')
  8. tensor([[1., 1., 1.],
  9. [1., 1., 1.]], device='cuda:1')

指定GPU计算神经网络模型

  1. from torch import nn
  2. net = nn.Linear(3, 1)
  3. print(list(net.parameters())[0].device) # cpu
  4. net.cuda(0)
  5. print(list(net.parameters())[0].device) # cuda:0
  6. net2 = nn.Linear(3, 1, device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
  7. print(list(net2.parameters())[0].device) # cuda:0
  8. net3 = nn.Linear(3, 1).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
  9. print(list(net3.parameters())[0].device) # cuda:0
  10. net3 = net3.cpu()
  11. print(list(net3.parameters())[0].device) # cpu

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/你好赵伟/article/detail/851467
推荐阅读
相关标签
  

闽ICP备14008679号