为什么Pytorch死活要用第0块显卡--我和pytorch的恩怨情仇

为什么Pytorch一定要用第0块显卡

更新 (2020.11.23)

Pytorch调用显卡的问题在1.0版本前比较明显,建议各位尽量使用1.0以后的版本。其次,对于超大模型的训练,有时会需要多显卡并行运算,使用方法指路另一篇博文 点这里

问题

在使用Pytorch进行训练的时候,有一个奇怪的问题是,无论怎么在代码里指定显卡,最终运行时始终会在第0块显卡上占用少量的内存。本来是个无伤大雅的问题,但是一旦0号显卡被自己或者别人占满了,就很抓狂?公用服务器的悲伤啊。
心痛

原因

查了一下在GitHub上有一个关于这个问题的讨论,Pytorch初始化问题,大概是pytorch在初始化的时候会默认在第0块显卡上进行,会占用一定的显存(在torch/cuda/init.py:110中,拷贝了部分代码过来)。这就导致,在第0块显卡空闲内存不多时,程序会反复报Runtime Error的错,可以通过如下的方法来达到完全不使用第0块显卡的目的。

class device(object):
    """Context-manager that changes the selected device.

    Arguments:
        idx (int): device index to select. It's a no-op if this argument
            is negative.
    """

    def __init__(self, idx):
        self.idx = idx
        self.prev_idx = -1

    def __enter__(self):
        if self.idx is -1:
            return
        _lazy_init()
        self.prev_idx = torch._C._cuda_getDevice()
        if self.prev_idx != self.idx:
            torch._C._cuda_setDevice(self.idx)

    def __exit__(self, *args):
        if self.prev_idx != self.idx:
            torch._C._cuda_setDevice(self.prev_idx)
        return False

解决方法

强制指定GPU的方法不少,但是不是每个都对pytorch有用,我在这里列出我知道的几个。

  1. 最好用的、绝对不会出错的device (适用于单卡调试,简单快捷,多卡训练的童靴使用方法2比较好)

    一般使用device时大多数用来直接指定显卡号,但是pytorch那个尿性在你制定前就用了第0块进行初始化,所以保险的方法是使用with语句,在with块下再写pytorch的调用,亲测有效。

with torch.cuda.device(1):
	pass
  1. 指令运行前设置环境变量 (适用于多卡训练)
    在输入训练指令前加上CUDA_VISIBLE_DEVICES=1,2,即设置环境变量,程序只“看得见”第1、2块显卡,自然就不需要在使用第0块显卡了。但是有一个问题是,在设置环境变量之后,第1、2块显卡在程序的眼里就变成了第0、1块显卡,在后续使用中需要注意序号的问题,不然会出现invalid device ordinal的错误。使用成本最小,绝对不会出现cuda0在被禁用的情况下被调用的情况。
CUDA_VISIBLE_DEVICES=1,2 python train.py
  1. 别人的博客上说有用,但是我自己没实验成功的os指定法
    用os指定理论上可以指定pytorch不走第0块GPU,但是我试了一下没成功,不知道是不是还要什么先决条件,有谁成功了欢迎评论指正。
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1,2'

Pytorch可真是个小妖精
狗脸一红在这里插入图片描述

  • 29
    点赞
  • 78
    收藏
    觉得还不错? 一键收藏
  • 11
    评论
评论 11
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值