为什么循环通过 pytorch 张量如此缓慢（与 Numpy 相比）？答案

【问题标题】：Why is looping through pytorch tensors so slow (compared to Numpy)?为什么循环通过 pytorch 张量如此缓慢（与 Numpy 相比）？
【发布时间】：2021-01-16 01:42:33
【问题描述】：

我最近一直在处理图像转换，遇到了一个情况，我有一个大数组（形状为 100,000 x 3），其中每一行代表 3D 空间中的一个点，例如：

pnt = [x y z]

我要做的只是遍历每个点和矩阵，将每个点与称为 T（形状 = 3 X 3）的矩阵相乘。

用 Numpy 测试：

def transform(pnt_cloud, T):
    
    i = 0
    for pnt in pnt_cloud:
        xyz_pnt = np.dot(T, pnt)
        
        if xyz_pnt[0] > 0:
            arr[i] = xyz_pnt[0]
            
        i += 1
           
    return arr

调用以下代码并计算运行时间（使用 %time）给出输出：

Out[190]: CPU times: user 670 ms, sys: 7.91 ms, total: 678 ms
Wall time: 674 ms

使用 Pytorch 张量进行测试：

import torch

tensor_cld = torch.tensor(pnt_cloud)
tensor_T   = torch.tensor(T)

def transform(pnt_cloud, T):
    depth_array = torch.tensor(np.zeros(pnt_cloud.shape[0]))

    i = 0
    for pnt in pnt_cloud:
        xyz_pnt = torch.matmul(T, pnt)
        
        if xyz_pnt[0] > 0:
            depth_array[i] = xyz_pnt[0]
            
        i += 1
            
        
    return depth_array

调用以下代码并计算运行时间（使用 %time）给出输出：

Out[199]: CPU times: user 6.15 s, sys: 28.1 ms, total: 6.18 s
Wall time: 6.09 s

注意：对 torch.jit 执行相同操作只会减少 2 秒

我原以为 PyTorch 张量计算会更快，因为 PyTorch 在编译阶段分解其代码的方式。我在这里错过了什么？

除了使用 Numba 之外，还有其他更快的方法吗？

【问题讨论】：

Quibble：如何将形状 (3,) 的点与形状 (4,4) 的矩阵相乘？尺寸不兼容吗？
pytorch 是否有可能为此操作累积梯度？
哦，是的，我删除了第四个术语以使问题更易于理解和通用，并且忘记减小另一个矩阵的大小。我已经编辑了问题，感谢您指出 Natchiket

标签： python numpy pytorch tensor

【解决方案1】：

为什么要使用 for 循环？
为什么你计算一个 3x3 的点积并且只使用结果的第一个元素？？

您可以在一个 matmul 中完成所有数学运算：

with torch.no_grad():
  depth_array = torch.matmul(pnt_cloud, T[:1, :].T)  # nx3 dot 3x1 -> nx1
  # since you only want non negative results
  depth_array = torch.maximum(depth_array, 0)

由于您想将运行时与 numpy 进行比较，您应该禁用 gradient accumulation。

【讨论】：

在上面的问题中，我声明 pnt_cld 是一个 100,000 x 3 矩阵，即每行是 1 x 3。我将每个 1 x 3 向量与 T (3x3) 相乘。 [nx3 dot 3x1 -> nx3 不等于上述过程]
我发现的另一件事是张量运算本身不会累积梯度。当操作与 nn.model 对象有关时，唯一的累积梯度。
@dankpenny 我更新了代码中关于矩阵大小的注释。由于您只使用结果的一个元素 - 您可以乘以矩阵的 3x1 行，而不是整个矩阵。
哦，是的，我的坏哈哈！

【解决方案2】：

为了速度，我从 PyTorch 论坛得到了这样的回复：

1-3 个元素的操作在 PyTorch 中通常相当昂贵，因为创建张量的开销变得很大（这包括设置单个元素），我认为这是这里的主要内容。这也是 JIT 没有太大帮助的原因（它只消除了 Python 开销）而 Numby 大放异彩（例如，对 depth_array[i] 的赋值只是内存写入）。
如果您在 PyTorch 和 NumPy 中有不同的 BLAS 后端，matmul 本身的速度可能会有所不同。

【讨论】：