Question

我正在计算大小为(50500,)的向量之间的巨大外部乘积，发现NumPy这样做比PyTorch快很多。

以下是测试：

# NumPy

In [64]: a = np.arange(50500) 
In [65]: b = a.copy()  

In [67]: %timeit np.outer(a, b) 
5.81 s ± 56.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

-------------

# PyTorch

In [73]: t1 = torch.arange(50500)
In [76]: t2 = t1.clone()

In [79]: %timeit torch.ger(t1, t2) 
7.73 s ± 143 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

理想情况下，我想在PyTorch中完成计算。那么，如何在PyTorch中针对如此巨大的向量加快计算外部乘积的速度呢？

注意：我试图将张量移动到GPU，但是我用MemoryError处理过，因为它需要大约19 GiB的空间。因此，我最终必须在CPU上完成该操作。

Answer 1

不幸的是，如果没有大量的努力，实际上是没有办法专门加快炬的外部产品torch.ger()的计算方法。

说明和选项

numpy函数np.outer()这么快的原因是因为它是用C编写的，您可以在这里看到：https://github.com/numpy/numpy/blob/7e3d558aeee5a8a5eae5ebb6aef03de892a92ebd/numpy/core/numeric.py#L1123 该函数使用来自umath C源代码的操作。

Pytorch的{{1}}函数是用C ++编写的：https://github.com/pytorch/pytorch/blob/7ce634ebc2943ff11d2ec727b7db83ab9758a6e0/aten/src/ATen/native/LinearAlgebra.cpp#L142，这使得它变慢了，如您在示例中看到的那样。

您“加速在PyTorch中计算外部产品的速度”的选项是在pytorch的本机代码中为外部产品添加C实现，或者在您确实不使用Cython与C进行接口的同时，使自己的外部产品功能运行。不想使用numpy（这没有多大意义）。

PS

此外，使用GPU只会提高您在GPU上的并行计算速度，可能不会超过在RAM和GPU内存之间传输数据所需的时间成本。

Answer 2

一个很好的解决方案是将两者结合起来。

class LazyFrames(object):
    def __init__(self, frames):
        self._frames = frames

    def __array__(self, dtype=None):
        out = np.concatenate(self._frames, axis=0)
        if dtype is not None:
            out = out.astype(dtype)
        return out

例如，

frames可能只是您的pytorch张量。

该对象确保观察之间的公共框架仅存储一次。它的存在纯粹是为了优化可能非常庞大的内存使用情况（例如DQN的1M帧重播缓冲区）。该对象只应在传递给模型之前转换为numpy数组。

参考：https://github.com/Shmuma/ptan/blob/master/ptan/common/wrappers.py

对于较大的交叉产品或外部产品，NumPy比PyTorch更快

2 个答案: