torch.cuda.comm.scatter

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

将张量分散到多个 GPU 上。

参数

• tensor (Tensor) – 要分散的张量。可以位于 CPU 或 GPU 上。

• devices (Iterable[torch.device, str 或 int], 可选) – 要分散到的 GPU 设备的迭代器。

• chunk_sizes (Iterable[int], 可选) – 要放置在每个设备上的块的大小。它在长度上应该与 devices 相匹配，并且应总计为 tensor.size(dim)。如果没有指定，tensor 将被分成相等大小的块。

• dim (int, 可选) – 沿其对 tensor 进行分块的维度。默认值：0。

• streams (Iterable[torch.cuda.Stream], 可选) – 要执行分散的流的迭代器。如果没有指定，将使用默认流。

• out (Sequence[Tensor], 可选, 关键字参数) – 用于存储输出结果的 GPU 张量。这些张量的大小必须与 tensor 相匹配，除了 dim，其中总大小必须总计为 tensor.size(dim)。

注意

必须指定 devices 或 out 中的其中一个。当指定 out 时，chunk_sizes 必须不指定，并将从 out 的大小推断。

返回值

• 如果指定了 devices，

  包含 tensor 块的元组，放置在 devices 上。

• 如果指定了 out，

  包含 out 张量的元组，每个张量都包含 tensor 的一个块。