当前位置 博文首页 > Pytorch如何切换 cpu和gpu的使用详解

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

    作者:遇到好事了 时间:2021-07-19 18:47

    前言,在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
    

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

    • 执行命令: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()
    jsjbwy
    下一篇:没有了