指定GPU训练的命令

一、单GPU训练

指定GPU：0 训练

CUDA_LAUNCH_BLOCKING=0 python train.py

指定其他GPU训练

CUDA_VISIBLE_DEVICES=1 python train.py

二、多GPU训练

需要修改代码

model = Your_model()
# 进行指定使用设备，这样会修改pytorch指定的设备编号如上则把device:2 改为 device:0，device:3 改为 device:1
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"  
# 则pytorch指定的编号还是从device:0 开始，指定使用程序为
model = torch.nn.DataParallel(model, device_ids=[0, 1]) 
model.to(device)

运行多GPU train脚本

python train_multi_GPU.py