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

169 阅读 0 评论 112 点赞

我是靠谱客的博主健康石头，这篇文章主要介绍pytorch使用过程中指定显卡训练，现在分享给大家，希望可以做个参考。

1 指定单块显卡进行训练

1.1 默认的0号显卡

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

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号卡进行训练