pytorch 设置多GPU

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

 

if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model, device_ids=[2,3])

如果出现 RuntimeError: all tensors must be on devices[0]：

该问题由于PyTorch程序设置的gpu_id编号未从0开始导致的。

解决方案，在运行python命令前设置一下CUDA_VISIBLE_DEVICES：

CUDA_VISIBLE_DEVICES=2,3 python example.py --gpu_id 0 1

注意：在程序中设置gpu_id仍要从0开始，上例中gpu_id 0对应device 2，gpu_id 1对应device 3