torch设置GPU

设置torch可见哪几个GPU：

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"  

# 例子>> "0"：只有0号GPU可见；"1,3"：只有1号和3号GPU可见；""：全部不可见，此时 torch.cuda.is_available() 输出 False。

注意：不一定非要写在最前面，但至少必须要在 import torch 之前，否则可能不起作用。

 

使用torch.cuda.set_device()可以更方便地将模型和数据加载到对应GPU上, 直接定义模型之前加入一行代码即可

torch.cuda.set_device(gpu_id) #单卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #可指定多卡

但是这种写法的优先级低，如果model.cuda()中指定了参数，那么torch.cuda.set_device()会失效，而且pytorch的官方文档中明确说明，不建议用户使用该方法。