Pytorch 多GPU运行

1、先设置当前环境可见的GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"]="2,3"#如希望使用2,3号GPU

2、创建设备

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")#使用上面设定的全部GPU
#device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#若用单块GPU，指定cuda:0时对应上面设定的2号GPU

3、将模型并行化

model=My_Model()
if torch.cuda.device_count()>1:#多CPU时
	model=nn.DataParallel(model)
model.to(device)

参考：
https://blog.csdn.net/xys430381_1/article/details/106635977?utm_medium=distribute.pc_relevant.none-task-blog-2_defaultbaidujs_title~default-0.control&spm=1001.2101.3001.4242