Как DataParallel выясняет, какой GPU я хочу использовать?

Question

Я хочу найти простой способ указать gpus, на котором выполняются мои эксперименты. В настоящее время я знаю, что могу использовать команду prepend моей команды python с CUDA_VISIBLE_DEVICES = 1,2,3,4 для установки графического процессора, и я предполагаю, что DataParallel затем попытается использовать все графические процессоры.

Есть ли способ сообщить DataParallel напрямую идентификаторы, например 4,7,9,12?

Answer 1

Да, DataParallem предоставляет нам возможность прямой передачи идентификаторов GPU.

Согласно официальной документации здесь , параллелизм данных реализован с использованием torch.nn.DataParallel.Модуль можно обернуть в DataParallel, и он будет распараллелен на несколько графических процессоров в пакетном измерении.

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

В вашем случае вы можете просто сделать что-то вроде этого:

torch.nn.DataParallel(model, device_ids=[4, 7, 9, 12])
output = net(input_var)  # input_var can be on any device, including CPU

Вы можете узнать больше о том, как передавать идентификаторы GPU напрямую в DataParallel по ссылкам ниже:

• ПРИМЕРЫ MULTI-GPU
• DataParallelслои

Answer 2

Используя DataParallel, вы можете указать, какие устройства вы хотите использовать с синтаксисом:

model = torch.nn.DataParallel(model, device_ids=[ids_1,ids_2, ..., ids_n]).cuda()

Когда вы используете CUDA_VISIBLE_DEVICES, вы устанавливаете графический процессор, видимый вашим кодом.Для удобства, если вы установите CUDA_VISIBLE_DEVICES = 2,3, а затем выполните:

model = torch.nn.DataParallel(model, device_ids=[0, 1]).cuda()

Используя nvidia-smi, вы увидите, что ваша модель будет работать на GPU 2 и 3, так как вы устанавливаете этот конкретный идентификатор каквиден.