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 or 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 匹配，除了 dim 維度，在該維度上，總大小必須等於 tensor.size(dim)。

注意

必須且只能指定 devices 和 out 中的一個。當指定 out 時，不得指定 chunk_sizes，塊大小將從 out 的大小推斷。

返回

• 如果指定了 devices，

  一個元組，包含 tensor 的塊，這些塊放置在 devices 上。

• 如果指定了 out，

  一個元組，包含 out 張量，每個張量都包含 tensor 的一個塊。