指定GPU 进行深度学习训练和多卡并行训练

指定GPU 进行深度学习训练

# 方法1：
torch.cuda.set_device(1)
# 方法2：
device = torch.device("cuda:1")
# 方法3：
# 推荐使用
os.environ["CUDA_VISIBLE_DEVICES"] = '1'

多GPU并行代码

使用四卡进行并行训练

    多GPU并行
    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
    if torch.cuda.device_count() > 1:
        model=torch.nn.DataParallel(model,device_ids=[0,1,2,3]).cuda() # 多GPU并行训练
    model.to(device)