torch 指定显卡

1. 代码中指定

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

2. shell中指定

export CUDA_VISIBLE_DEVICES=gpu_ids && python3 train.py

3. 模型/参数中指定

model.cuda(gpu_id) # gpu_id为int类型变量，只能指定一张显卡
model.cuda('cuda:'+str(gpu_ids)) #输入参数为str类型，可指定多张显卡
model.cuda('cuda:1,2') #指定多张显卡的一个示例

model.cuda(1)
loss.cuda(1)
tensor.cuda(1)

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

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