189 8069 5689

怎么在Pytorch中切换cpu和gpu-创新互联

本篇文章给大家分享的是有关怎么在Pytorch中切换 cpu和gpu,小编觉得挺实用的,因此分享给大家学习,希望大家阅读完这篇文章后可以有所收获,话不多说,跟着小编一起来看看吧。

创新互联公司不只是一家网站建设的网络公司;我们对营销、技术、服务都有自己独特见解,公司采取“创意+综合+营销”一体化的方式为您提供更专业的服务!我们经历的每一步也许不一定是最完美的,但每一步都有值得深思的意义。我们珍视每一份信任,关注我们的成都做网站、成都网站制作、成都外贸网站建设质量和服务品质,在得到用户满意的同时,也能得到同行业的专业认可,能够为行业创新发展助力。未来将继续专注于技术创新,服务升级,满足企业一站式全网营销推广需求,让再小的成都品牌网站建设也能产生价值!

方法1:x.to(device)


把 device 作为一个可变参数,推荐使用argparse进行加载:

使用gpu时:


device='cuda'
x.to(device) # x是一个tensor,传到cuda上去

使用cpu时:


device='cpu'
x.to(device)

方法2:使用x.cuda()+CUDA_VISIBLE_DEVICES

很多贴子中说,使用x.cuda() 和x.to('cuda') 虽然是等效的,但是x.cuda() 的缺点是无法动态切换cpu。然而,其实配合命令行参数CUDA_VISIBLE_DEVICES 是可以进行切换的。

在服务器上创建一个python脚本 t.py:

import torch
print(torch.cuda.device_count()) # 可用gpu数量
print(torch.cuda.is_available()) # 是否可用gpu

首先先看一下,正常运行的情况:

  • 执行命令:python t.py

  • 输出结果:因为服务器上有两个gpu,所以是我们想要的结果。

2
True


如果想要只使用某一块gpu,只需要在执行前加一个参数:

  • CUDA_VISIBLE_DEVICES=0 python t.py,例如,我们要使用gpu 0

  • 接下来看看输出什么:是的!程序中确实只可见了一块gpu~


1
True


下面,如果我们想使用cpu呢?

  • CUDA_VISIBLE_DEVICES="" python t.py

  • 输出结果:可以看到,虽然服务器上有2块cpu,通过我们设置执行参数,程序中也成功看不到了!


0
False


因此,回归正题,当我们使用x.cuda()进行分配gpu时,只需要使用torch.cuda.is_available()加一个判断即可,当想使用cpu的时候在执行程序的命令行参数进行控制:

if torch.cuda.is_available():
  x= x.cuda()

以上就是怎么在Pytorch中切换 cpu和gpu,小编相信有部分知识点可能是我们日常工作会见到或用到的。希望你能通过这篇文章学到更多知识。更多详情敬请关注创新互联行业资讯频道。


本文标题:怎么在Pytorch中切换cpu和gpu-创新互联
本文地址:http://cdxtjz.cn/article/iphpp.html

其他资讯