torch.cuda.comm.scatter

torch.cuda.comm.scatter(tensor, devices=None, chunk_sizes=None, dim=0, streams=None, *, out=None) [source]

Разделение тензора по нескольким GPU.

Параметры:

• tensor (Тензор) – тензор для разделения. Может находиться на процессоре или GPU.
• devices (Iterable[torch.device, str или int], необязательно) – итерируемый набор устройств GPU, на которых необходимо выполнить разделение.
• chunk_sizes (Iterable[int], необязательно) – размеры фрагментов, которые будут размещены на каждом устройстве. Длина должна соответствовать devices, а сумма размеров равна tensor.size(dim). Если не указано, то tensor будет разделен на равные фрагменты.
• dim (int, необязательно) – размерность, по которой разбить tensor. Значение по умолчанию: 0.
• streams (Iterable[Поток], необязательно) – итерируемый набор потоков, в рамках которых выполняется разделение. Если не указано, используется пото́к по умолчанию.
• out (Sequence[Тензор], необязательно, только ключевое слово) – тензоры GPU для хранения результатов. Размеры этих тензоров должны соответствовать tensor, за исключением dim, где общий размер должен быть равен tensor.size(dim).

Примечание

Должен быть указан ровно один из devices и out . При указании out, chunk_sizes не должен быть указан и будет выведен из размеров out.

Возвращает:

• If devices is specified,

  кортеж, содержащий фрагменты tensor, размещённые на devices.

• If out is specified,

  кортеж, содержащий тензоры out, каждый из которых содержит фрагмент tensor.