pytorch使用过程中指定显卡训练

目录


1 指定单块显卡进行训练

        1.1 默认的0号显卡

        若使用“0”号显卡进行训练,可以直接在代码中写

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu' )

        1.2 使用其他显卡,例如1号或其他

        若直接将1.1中的cuda:0改为cuda:1,则加载模型和数据还是通过0号显卡,如果0号卡被占满,那么将导致报“cuda out of memory”的问题。

        如果需要所有过程都通过1号显卡,那么可以直接用os指定。代码如下

import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )

# 接下来的代码可能是
# net.to(device)
# 为了把模型和数据都转移到这个device中,由于限定了可见cuda为1,故能指定只使用1号卡进行训练

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
乘风的头像乘风管理团队
上一篇 2023年2月28日 下午10:06
下一篇 2023年2月28日

相关推荐