torch.cuda.comm.scatter

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

在多个GPU上分散张量。

Parameters

• 张量 (张量) – 要分散的张量。可以在CPU或GPU上。

• 设备（可迭代[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 (序列[Tensor], 可选, 仅关键字) – 用于存储输出结果的GPU张量。这些张量的大小必须与tensor的大小匹配，除了dim维度外，该维度的总大小必须等于tensor.size(dim)。

注意

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

Returns

• 如果指定了devices，

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

• 如果指定了out，

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