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) – 用於執行 scatter 的 Stream 可迭代物件。如果未指定，將使用預設的 stream。

• out (Sequence[Tensor], optional, keyword-only) – 用於儲存輸出結果的 GPU tensors。這些 tensors 的大小必須與 tensor 的大小相符，除了 dim 維度，其總大小必須等於 tensor.size(dim)。

注意

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

回傳值

• 如果指定了 devices，

  一個包含 tensor 的區塊的 tuple，放置在 devices 上。

• 如果指定了 out，

  一個包含 out tensors 的 tuple，每個 tensor 包含 tensor 的一個區塊。