时间戳

首页 > 解决方案 > 为什么循环通过 pytorch 张量如此缓慢（与 Numpy 相比）？

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

问题描述

我最近一直在处理图像转换，遇到了一种情况，我有一个大数组（形状为 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 之外，还有其他更快的方法吗？

标签： pythonnumpypytorchtensor

解决方案

你为什么使用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 进行比较，因此应禁用梯度累积。

推荐阅读