torch.cuda.comm.scatter

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

将 tensor 分散到多个 GPU 中。

参数

• tensor (Tensor) – 要分散的 tensor。可以是 CPU 或 GPU 上的 tensor。

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

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

• dim (int, optional) – 要将 tensor 分块的维度。默认值：0。

• streams (Iterable[torch.cuda.Stream], optional) – 要在其上执行分散操作的 Stream 的可迭代对象。如果未指定，将使用默认 Stream。

• out (Sequence[Tensor], optional, keyword-only) – 用于存储输出结果的 GPU tensor。这些 tensor 的大小必须与 tensor 匹配，除了 dim 维度，该维度的总大小必须等于 tensor.size(dim)。

注意

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

返回

• 如果指定了 devices，

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

• 如果指定了 out，

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