Загружает ли PyTorch nn.DataParallel одну и ту же модель в каждый графический процессор?

Question

Единственный способ, которым это могло бы работать (логически), - это если бы модель была загружена в каждый из графических процессоров. Это будет означать, что при обновлении весов каждому графическому процессору также потребуется обновить веса, увеличивая рабочую нагрузку по сравнению с одним графическим процессором. Верна ли эта аргументация?

Answer 1

Прежде всего, рекомендуется использовать вместо torch.nn.parallel.DistributedDataParallel.

Вы можете проверить torch.nn.DataParallel документацию, где описан процесс (вы также можете проверить исходный код и немного глубже покопаться в github , здесь как выполняется репликация модуля).

Вот примерно то, как это делается:

Инициализация

Все (или выбранные) идентификаторы устройств сохраняются в конструкторе и измерении, по которым будут разбросаны данные (почти всегда 0 означает быть разделено на устройства по пакету)

Вперед

Это делается во время каждого forward прогона:

1. Входы разбросаны (тензоры по размерам, tuple , list, dict неглубоко скопировано, другие данные распределяются между потоками).
2. Если есть только одно устройство, просто верните module(*args, **kwargs)
3. Если есть несколько устройств, скопируйте сеть с исходной машины на другие устройства (это делается каждый раз!)
4. Вперед на каждом устройстве с соответствующим входом
5. Собирать выходные данные с устройств на одно исходное устройство (конкатенация выходов) на исходный компьютер.
6. Выполнить остальную часть кода, выполнить обратную передачу, обновить веса на Исходная машина et c.

Исходная машина - это cuda:0 по умолчанию, но ее можно выбрать. Также веса обновляются для одного device, только пакет разбрасывается, а выходы собираются.