亲宝软件园·资讯

展开

Pytorch切换cpu和gpu Pytorch怎样切换 cpu和gpu的使用详解

遇到好事了 人气:0
想了解Pytorch怎样切换 cpu和gpu的使用详解的相关内容吗,遇到好事了在本文为您仔细讲解Pytorch切换cpu和gpu的相关知识和一些Code实例,欢迎阅读和指正,我们先划重点:Pytorch切换cpu和gpu,Pytorch,cpu,gpu,下面大家一起来学习吧。

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

方法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

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

2
True

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

1
True

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

0
False

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

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

加载全部内容

相关教程
猜你喜欢
用户评论