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 的一個塊。