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], 可选) – 要在其上执行分散的 Stream 的可迭代对象。如果未指定，将使用默认流。

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

注意

必须指定 devices 和 out 中的一个。当指定 out 时，不得指定 chunk_sizes，并且将从 out 的大小推断出来。

返回

• 如果指定了 devices，

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

• 如果指定了 out，

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