Использование `DataParallel`, когда сети нужен общий (постоянный)` Tensor`

Question

Я хотел бы использовать DataParallel для распределения моих вычислений по нескольким графическим процессорам по пакетному измерению. Моя сеть требует Tensor (назовем это A) внутри, что является постоянным и не меняется при оптимизации. Кажется, что DataParallel не копирует автоматически этот Tensor во все рассматриваемые графические процессоры, и сеть, таким образом, будет жаловаться, что фрагмент входных данных x, которые он видит, находится на другом графическом процессоре, чем A.

Есть ли способ DataParallel автоматически справиться с этой ситуацией? В качестве альтернативы, есть ли способ скопировать Tensor на все графические процессоры? Или мне просто оставить один Tensor для каждого графического процессора и вручную выяснить, какую копию использовать, в зависимости от того, где находится фрагмент, видимый forward?

Answer 1

Вы должны обернуть свой тензор в torch.nn.Parameter и установить requires_grad=False во время его создания.

torch.nn.Parameter не означает, что тензор должен быть обучаемым .

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

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

Я не вижу необходимости в другой функции, чтобы сделать это, хотя имя может немного сбить с толку.