服务器之家

服务器之家 > 正文

Pytorch如何切换 cpu和gpu的使用详解

时间:2021-09-12 00:49     来源/作者:遇到好事了

前言,在pytorch中,当服务器上的gpu被占用时,很多时候我们想先用cpu调试下代码,那么就需要进行gpu和cpu的切换。

方法1:x.to(device)

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

使用gpu时:

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

使用cpu时:

?
1
2
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:

?
1
2
3
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的时候在执行程序的命令行参数进行控制:

?
1
2
if torch.cuda.is_available():
  x= x.cuda()

到此这篇关于pytorch如何切换 cpu和gpu的使用详解的文章就介绍到这了,更多相关pytorch切换cpu和gpu内容请搜索服务器之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持服务器之家!

原文链接:https://blog.csdn.net/qq_42549774/article/details/114192405

标签:

相关文章

热门资讯

yue是什么意思 网络流行语yue了是什么梗
yue是什么意思 网络流行语yue了是什么梗 2020-10-11
2020微信伤感网名听哭了 让对方看到心疼的伤感网名大全
2020微信伤感网名听哭了 让对方看到心疼的伤感网名大全 2019-12-26
背刺什么意思 网络词语背刺是什么梗
背刺什么意思 网络词语背刺是什么梗 2020-05-22
苹果12mini价格表官网报价 iPhone12mini全版本价格汇总
苹果12mini价格表官网报价 iPhone12mini全版本价格汇总 2020-11-13
2021年耽改剧名单 2021要播出的59部耽改剧列表
2021年耽改剧名单 2021要播出的59部耽改剧列表 2021-03-05
返回顶部