torch.cuda.comm.scatter

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

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

Параметры:

• tensor (Tensor) – тензор для распределения. Может находиться на CPU или GPU.
• devices (Iterable[torch.device, str or int], optional) – итерируемый объект с устройствами GPU, между которыми распределяется тензор.
• chunk_sizes (Iterable[int], optional) – размеры фрагментов, размещаемых на каждом устройстве. Длина должна совпадать с devices, а сумма — равняться tensor.size(dim). Если параметр не указан, tensor будет разделён на равные фрагменты.
• dim (int, optional) – измерение, по которому разбивается tensor. По умолчанию: 0.
• streams (Iterable[torch.cuda.Stream], optional) – итерируемый объект с потоками, в которых выполняется распределение. Если параметр не указан, будет использоваться поток по умолчанию.
• out (Sequence[Tensor], optional, keyword-only) – тензоры GPU для сохранения результатов. Размеры этих тензоров должны совпадать с размером tensor, за исключением dim, для которого общий размер должен равняться tensor.size(dim).

Примечание

Необходимо указать ровно один из параметров devices и out. Если указан out, параметр chunk_sizes указывать нельзя; его значение будет определено по размерам out.

Возвращает:

• If devices is specified,

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

• If out is specified,

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