交流集

torch.cuda.comm.broadcast(tensor, devices)

向一些GPU广播张量。

参数：

• tensor (Tensor) – 将要广播的张量
• devices (Iterable) – 一个可以广播的设备的迭代。注意，它的形式应该像（src，dst1，dst2，...），其第一个元素是广播来源的设备。

返回： 一个包含张量副本的元组，放置在与设备的索引相对应的设备上。

torch.cuda.comm.reduce_add(inputs, destination=None)

将来自多个GPU的张量相加。

所有输入应具有匹配的形状。

参数：

• inputs (Iterable[Tensor]) – 要相加张量的迭代
• destination (int, optional) – 将放置输出的设备（默认值：当前设备）。

返回： 一个包含放置在destination设备上的所有输入的元素总和的张量。

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

打散横跨多个GPU的张量。

参数：

• tensor (Tensor) – 要分散的张量
• devices (Iterable[int]) – int的迭代，指定哪些设备应该分散张量。
• chunk_sizes (Iterable[int], optional) – 要放置在每个设备上的块大小。它应该匹配devices的长度并且总和为tensor.size(dim)。 如果没有指定，张量将被分成相等的块。
• dim (int, optional) – 沿着这个维度来chunk张量

返回： 包含tensor块的元组，分布在给定的devices上。

torch.cuda.comm.gather(tensors, dim=0, destination=None)

从多个GPU收集张量。

张量尺寸在不同于dim的所有维度上都应该匹配。

参数：

• tensors (Iterable[Tensor]) – 要收集的张量的迭代。
• dim (int) – 沿着此维度张量将被连接。
• destination (int, optional) – 输出设备（-1表示CPU，默认值：当前设备）。

返回： 一个张量位于destination设备上，这是沿着dim连接tensors的结果。