SpConv 库中体素生成逻辑的 Numpy 实现以及性能测试

已于 2025-03-28 17:57:00 修改
阅读量694 收藏 10
点赞数 17
文章标签： numpy
于 2025-03-28 11:31:32 首次发布
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.youkuaiyun.com/m0_62121252/article/details/146590115
版权

SpConv 库中体素生成逻辑的 Numpy 实现以及性能测试

起因

因为实验室服务器重新分配的缘故, 我不得不将实验环境迁移到老师提供的其他服务器上. 然而当我重新配置好代码环境之后, 准备开跑, 发现代码无法运行, 在加载数据的时候会出现浮点数运算错误.

经过查询可知, 这是因为 SpConv 库与 Numpy2.0 版本不兼容所致, 将代码降级到 Numpy2.0 以下就可以解决这个问题.

然后这件事情一直让我耿耿于怀, 总觉得不能用上最新版的 Numpy 库是一个大问题. 因为鄙人之前还浮现过其他论文的实验, 因为那些论文都比较古早, 不仅 pytorch 版本落后, 就算是 python 版本也早已跟不上时代,

考虑到未来我的许多工作都要基于这个论文的代码进行修改, 不能因为这个一个 SpConv 库导致我不能使用 Numpy2.0 (虽然我也不清楚这个 2.0 相比起 1.x 到底有哪些方面的变化 (重点是在性能方面的变化), 但是升级总是比不升级好的 (至少我是这么想的)). 考虑到我只是使用了 SpConv 中体素生成的接口, 并没有使用其中更复杂的什么稀疏卷积之类的, 所以实现起来应该比较简单.

需求分析

因为鄙人所以依赖的项目是 HEAL, 使用的数据集是 OPV2V. 所以我对这次仿写指定的目标是:

  • 查看 HEAL 代码中使用到 SpConv 库接口的名称以及传参格式
  • 查看 SpConv 的源代码, 找出这个接口对应的实现代码
  • 对源代码进行分析, 同时实现
  • 在 OPV2V 数据集上验证我代码的正确性, 然后做一下性能比较

具体实现过程

查看 HEAL 代码中使用到 SpConv 库接口的名称以及传参格式

在数据处理部分只有生成体素的时候使用到了 SpConv 库中生成体素的接口, 具体的文件地址: opencood/data_utils/pre_processor/sp_voxel_preprocessor.py. 以及两张代码截图:
初始化时 SpConv生成体素的用发
可以看到 HEAL 中同时支持 SpConv 1.x 和 SpConv 2.x. 因为鄙人使用的时 SpConv 2.x 所以我们具体看 2.x 部分使用的代码, 这部分代码我已经用红框进行了圈画. 通过分析我们不难知道 HEAL 使用了 SpConv 这个接口 :

from spconv.utils import Point2VoxelCPU3d as VoxelGenerator

这样只要我们去 SpConv 源代码下搜索这个库就能顺利找到他对应的实现了, 然后就可以愉快地进行仿写了

查看 SpConv 的源代码, 找出这个接口对应的实现代码, 然后进行仿写

然而, SpConv 中并没有这么一个相同的类名. 因为 SpConv 为了支持多种不同的维度, 只写了一个实现: Point2VoxelCPU (路径在 spconv/spconv/csrc/sparse/pointops.py) 然后同目录下的 all.py 再根据维度的不同生成维度的实现, 最后通过 utils 文件夹下的 __inti__.py 暴露出去.

# utils/__init__.py 中的引入逻辑, 从中可以看到 all.py 的生成逻辑
from spconv.core_cc.csrc.sparse.all.ops_cpu1d import Point2VoxelCPU as Point2VoxelCPU1d
from spconv.core_cc.csrc.sparse.all.ops_cpu2d import Point2VoxelCPU as Point2VoxelCPU2d
from spconv.core_cc.csrc.sparse.all.ops_cpu3d import Point2VoxelCPU as Point2VoxelCPU3d
from spconv.core_cc.csrc.sparse.all.ops_cpu4d import Point2VoxelCPU as Point2VoxelCPU4d

但以上这些都不是重点, 所以我也不详细讲, 总结一句话: 因为不同维度的 point 生成 voxel 的逻辑基本一致, 所以为了省事, SpConv 没有写 4 份代码, 即使写了 4 份代码, 也就是维度那个地方改一改, 其他的什么也没有变. 所以 SpConv 就只实现了一个代码模板, 然后传入不同的 ndim 生成这 4 份代码. 要想重写这个生成体素的代码, 我们只需要关注 Point2VoxelCPU 这个类就可以了.

对源代码进行分析, 同时实现

通过强大的 Python IDE PyCharm, 我们可以看到 Point2VoxelCPU 这个类的结构:
Point2VoxelCPU 类结构
重点关注的函数是 ctor(self) , point_to_voxel(self), point_to_voxel_static(self) 以及 point_to_voxel_static_template(self, mean: bool = False) 而具体要分析的代码就在最后一个函数中. 下面我们就结合代码仔细看看这个函数

这段代码是声明生成的 C++ 函数的参数 (我通过问 Grok 得到的)

code = pccm.FunctionCode()
code.arg("points", "tv::Tensor")
code.arg(
    "voxels, indices, num_per_voxel, densehashdata, points_voxel_id", "tv::Tensor"
)
code.arg("vsize", f"std::array<float, {self.ndim}>")
code.arg("grid_size, grid_stride", f"std::array<int, {self.ndim}>")
code.arg("coors_range", f"std::array<float, {self.ndim * 2}>")

code.arg("clear_voxels", "bool", "true")

这段代码呢, 则是下面主要逻辑部分循环时用到的变量 (根据点的顺序, 进行一下字符串的替换)

point_xyz = f"{self.ndim - 1} - j"
if not self.zyx:
    point_xyz = f"j"

下面的代码是主体逻辑, 注释写到代码里面了, 同时省略了最外层的 code.raw( f"""

auto max_num_voxels = voxels.dim(0);
auto max_num_points_per_voxel = voxels.dim(1);
num_per_voxel.zero_();
if (clear_voxels) {{
    voxels.zero_();
}}
auto points_voxel_id_ptr = points_voxel_id.data_ptr<int64_t>();
int res_voxel_num = 0;
int num_features = points.dim(1);
auto N = points.dim(0);
int c;
// 检查传参, 传入的点云可能不只有 x, y, 
TV_ASSERT_RT_ERR(num_features == voxels.dim(2), "your points num features doesn't equal to voxel.");
// Grok: dispatch 是一个函数模板, 用于根据输入的数据类型动态分发 (dispatch) 到不同的处理逻辑
// Grok: <float, double> 是模板参数, 表示 dispatch 支持的类型列表, 这里限制为 float 和 double
tv::dispatch<float, double>(points.dtype(), [&](auto I){{
    using T = decltype(I);
    // Grok: tview<T, 2> 表示将 points（可能是一个张量或多维数组）转换为一个二维视图, 并指定数据类型为 T (这里 T 是通过 decltype(I) 从 dispatch 中推导出来的，比如 float 或 double).
    // 这些变量的定义都可以在 `ctor` 函数中找到
    auto points_rw = points.tview<T, 2>();
    auto coors_rw = indices.tview<int, 2>();
    auto voxels_rw = voxels.tview<{self.dtype}, 3>();
    auto num_points_per_voxel_rw = num_per_voxel.tview<int, 1>();
    
    // 单个点体素的坐标
    int coor[{self.ndim}];
    auto coor_to_voxelidx_rw = densehashdata.tview<int, {self.ndim}>();
    int voxelidx, num;
    bool failed;
    int voxel_num = 0;
    for (int i = 0; i < N; ++i) {{
        failed = false;
        for (int j = 0; j < {self.ndim}; ++j) {{
            // 这个地方原来就是作了一个简单的字符串替换
            c = floor((points_rw(i, {point_xyz}) - coors_range[j]) / vsize[j]);
            if ((c < 0 || c >= grid_size[j])) {{
                failed = true;
                break;
            }}
            coor[j] = c;
        }}
        if (failed){{
            // 给这个 point 的体素的 id 设置为 -1
            points_voxel_id_ptr[i] = -1;
            continue;
        }}
        // Grok: 占位展开. 等价代码: voxelidx = coor_to_voxelidx_rw(coor[0], coor[1], coor[2])
        // 根据单个点的体素的坐标, 计算出体素的 id
        voxelidx = coor_to_voxelidx_rw({codeops.unpack("coor", range(self.ndim))});
        
        if (voxelidx == -1) {{
            // 如果体素 id 为 -1, 说明当前体素还没有分配 id, 则给这个体素分配一个 id (也就是当前的 voxel_num)
            // 之后, 根据体素的坐标把这个 id 存放到字典里面 (应该有两个字典的, C++ 实现使用的是数组)
            voxelidx = voxel_num;
            if (voxel_num >= max_num_voxels){{
                points_voxel_id_ptr[i] = -1;
                continue;
            }}
            voxel_num += 1;
            coor_to_voxelidx_rw({codeops.unpack("coor", range(self.ndim))}) = voxelidx;
            for (int k = 0; k < {self.ndim}; ++k) {{
                coors_rw(voxelidx, k) = coor[k];
            }}
        }}
        // 处理完体素 id 的问题之后
        points_voxel_id_ptr[i] = voxelidx;
        num = num_points_per_voxel_rw(voxelidx);
        if (num < max_num_points_per_voxel) {{
            // voxel_point_mask_rw(voxelidx, num) = {self.dtype}(1);
            for (int k = 0; k < num_features; ++k) {{
                voxels_rw(voxelidx, num, k) = points_rw(i, k);
            }}
            num_points_per_voxel_rw(voxelidx) += 1;
        }}
    }}
    // 这一段应该在求平均值
    std::vector<{self.dtype}> mean_value(num_features);
    for (int i = 0; i < voxel_num; ++i) {{
        coor_to_voxelidx_rw({codeops.unpack("coors_rw", range(self.ndim), left="(i, ", right=")")}) = -1;
        if TV_IF_CONSTEXPR ({pccm.boolean(mean)}){{
            num = num_points_per_voxel_rw(i);
            if (num > 0){{
                mean_value.clear();
                for (int j = 0; j < num; ++j) {{
                    for (int k = 0; k < num_features; ++k) {{
                        mean_value[k] += voxels_rw(i, j, k);
                    }}
                }}
                for (int k = 0; k < num_features; ++k){{
                    mean_value[k] /= num;
                }}
                for (int j = num; j < max_num_points_per_voxel; ++j) {{
                    for (int k = 0; k < num_features; ++k) {{
                        voxels_rw(i, j, k) = mean_value[k];
                    }}
                }}
            }}
        }}
    }}
    res_voxel_num = voxel_num;
}});
return std::make_tuple(voxels.slice_first_axis(0, res_voxel_num), 
    indices.slice_first_axis(0, res_voxel_num), 
    num_per_voxel.slice_first_axis(0, res_voxel_num));

关于里面 tensor (应该不算是 pytorch 里面的那种 tensor, 是 cumm 这个库里面的 tensor) 的 shape, 可以在 ctor 函数中找到, 或者可以结合我写的 Numpy 形式的代码进行推理, 在此不做过多追述, 下面直接开始写代码.

我一共写了三版代码. 第一版是完全按照该函数仿写的, 另外一版则是先过滤出在体素范围内的点, 然后对体素范围内的点分配体素 id 的代码, 还有一版是使用了 jit 注解的

import numpy as np
import numba as nb


@nb.jit(cache=True, nopython=True)
def point2voxel_jit(
    ndim: int,
    dtype: np.dtype,
    vsize: np.ndarray,
    coors_range: np.ndarray,
    num_point_features: int,
    max_num_voxels: int,
    max_num_points_per_voxel: int,
    grid_size: np.ndarray,
    # 以上是 self 所需要的参数
    points: np.ndarray,
):
    N: int = points.shape[0]

    points_voxel_id = np.full(N, -1, dtype=np.int32)
    voxels = np.full((max_num_voxels, max_num_points_per_voxel, num_point_features), -1, dtype=dtype)
    num_per_voxel = np.zeros(max_num_voxels, dtype=np.int32)

    # 两个双向的 map
    voxelidx_to_coor = np.full((max_num_voxels, ndim), -1, dtype=np.int32)
    # 这里把 tuple 拆开了
    coor_to_voxelidx = np.full((grid_size[0], grid_size[1], grid_size[2]), -1, dtype=np.int32)

    voxel_num = 0
    for i in range(N):
        failed = False
        coor = np.zeros(ndim, dtype=np.int32)

        """计算当前点对应的体素的坐标 (其实这里可以用一下 numpy 的, 而不用循环)"""
        for j in range(ndim):
            c = np.floor((points[i, j] - coors_range[j]) / vsize[j])
            if c < 0 or c >= grid_size[j]:
                failed = True
                break
            coor[j] = int(c)
        if failed:
            # 其实这里应该不用赋值的, 因为初始化的时候已经赋值为 -1 了
            points_voxel_id[i] = -1
            continue
        # 这里把 tuple 拆开了
        coor_tuple = (coor[0], coor[1], coor[2])
        voxelidx = coor_to_voxelidx[coor_tuple]
        if voxelidx == -1:
            voxelidx = voxel_num
            if voxel_num >= max_num_voxels:
                points_voxel_id[i] = -1
                continue
            voxel_num += 1
            coor_to_voxelidx[coor_tuple] = voxelidx
            voxelidx_to_coor[voxelidx] = coor

        points_voxel_id[i] = voxelidx
        num = num_per_voxel[voxelidx]
        if num < max_num_points_per_voxel:
            voxels[voxelidx, num] = points[i]
            num_per_voxel[voxelidx] += 1

    """TODO: 求平均值那部分就不抄写过来了"""
    return voxels[:voxel_num], voxelidx_to_coor[:voxel_num], num_per_voxel[:voxel_num]


class Point2VoxelNumpy3D:
    ndim = 3  # 要给他声明成一个不可变的东西
    dtype = np.float32

    def __init__(
        self,
        vsize_xyz: np.ndarray,
        coors_range_xyz: np.ndarray,
        num_point_features: int,
        max_num_voxels: int,
        max_num_points_per_voxel: int,
    ):
        """
        :param vsize_xyz: 体素的大小, shape: (3,)
        :param coors_range_xyz: 体素的范围, shape: (6,) - <min_x>, <min_y>, <min_z>, <max_x>, <max_y>, <max_z>
        """
        if vsize_xyz.ndim != self.ndim:
            Exception(f"vsize_xyz.ndim must be {self.ndim}.")
        if coors_range_xyz.ndim != 2 * self.ndim:
            Exception(f"coors_range_xyz.ndim must be {2 * self.ndim}.")
        self.vsize = vsize_xyz.astype(self.dtype)
        self.coors_range = coors_range_xyz.astype(self.dtype)
        self.num_point_features = num_point_features
        self.max_num_voxels = max_num_voxels
        self.max_num_points_per_voxel = max_num_points_per_voxel

        self.grid_size = np.floor((self.coors_range[3:] - self.coors_range[:3]) / self.vsize).astype(np.int32)

    def point2voxel(self, points: np.ndarray):
        """
        :param points: 传入的点云数据, shape: (N, num_point_features)
        :param mean: 是否进行平均 (这个参数为看 spcov 里面有, 但是具体到 HEAL 上没有用到, 所以就不写这部分的代码了)
        """
        if points.shape[1] != self.num_point_features:
            Exception(f"{points.shape[1]=} must be equal to {self.num_point_features=}.")
        N: int = points.shape[0]

        points_voxel_id = np.full(N, -1, dtype=np.int32)
        voxels = np.full((self.max_num_voxels, self.max_num_points_per_voxel, self.num_point_features), -1, dtype=self.dtype)
        num_per_voxel = np.zeros(self.max_num_voxels, dtype=np.int32)

        # 两个双向的 map
        voxelidx_to_coor = np.full((self.max_num_voxels, self.ndim), -1, dtype=np.int32)
        coor_to_voxelidx = np.full(tuple(self.grid_size), -1, dtype=np.int32)

        voxel_num = 0
        for i in range(N):
            failed = False
            coor = np.zeros(self.ndim, dtype=np.int32)

            """计算当前点对应的体素的坐标 (其实这里可以用一下 numpy 的, 而不用循环)"""
            for j in range(self.ndim):
                c = np.floor((points[i, j] - self.coors_range[j]) / self.vsize[j])
                if c < 0 or c >= self.grid_size[j]:
                    failed = True
                    break
                coor[j] = int(c)
            if failed:
                # 其实这里应该不用赋值的, 因为初始化的时候已经赋值为 -1 了
                points_voxel_id[i] = -1
                continue
            coor_tuple = tuple(coor)
            voxelidx = coor_to_voxelidx[coor_tuple]
            if voxelidx == -1:
                voxelidx = voxel_num
                if voxel_num >= self.max_num_voxels:
                    points_voxel_id[i] = -1
                    continue
                voxel_num += 1
                coor_to_voxelidx[coor_tuple] = voxelidx
                voxelidx_to_coor[voxelidx] = coor

            points_voxel_id[i] = voxelidx
            num = num_per_voxel[voxelidx]
            if num < self.max_num_points_per_voxel:
                voxels[voxelidx, num] = points[i]
                num_per_voxel[voxelidx] += 1

        """TODO: 求平均值那部分就不抄写过来了"""
        return voxels[:voxel_num], voxelidx_to_coor[:voxel_num], num_per_voxel[:voxel_num]

    def point2voxel2(self, points: np.ndarray):
        if points.shape[1] != self.num_point_features:
            Exception(f"{points.shape[1]=} must be equal to {self.num_point_features=}.")

        """筛选出在范围内的点, 并计算点的体素坐标"""
        points_coords = np.floor((points[:, :3] - self.coors_range[:3]) / self.vsize).astype(np.int32)
        mask = ((points_coords >= 0) & (points_coords < self.grid_size)).all(1)
        points, points_coords = points[mask], points_coords[mask, ::-1]

        N: int = points.shape[0]

        points_voxel_id = np.full(N, -1, dtype=np.int32)
        voxels = np.full((self.max_num_voxels, self.max_num_points_per_voxel, self.num_point_features), -1, dtype=self.dtype)
        num_per_voxel = np.zeros(self.max_num_voxels, dtype=np.int32)

        # 两个双向的 map
        voxelidx_to_coor = np.full((self.max_num_voxels, self.ndim), -1, dtype=np.int32)
        coor_to_voxelidx = np.full(tuple(self.grid_size), -1, dtype=np.int32)

        voxel_num = 0
        for i in range(N):
            coor = points_coords[i]

            voxelidx = coor_to_voxelidx[tuple(coor)]
            if voxelidx == -1:
                voxelidx = voxel_num
                if voxel_num >= self.max_num_voxels:
                    points_voxel_id[i] = -1
                    continue
                voxel_num += 1
                coor_to_voxelidx[tuple(coor)] = voxelidx
                voxelidx_to_coor[voxelidx] = coor

            points_voxel_id[i] = voxelidx
            num = num_per_voxel[voxelidx]
            if num < self.max_num_points_per_voxel:
                voxels[voxelidx, num] = points[i]
                num_per_voxel[voxelidx] += 1

        """TODO: 求平均值那部分就不抄写过来了"""
        return voxels[:voxel_num], voxelidx_to_coor[:voxel_num], num_per_voxel[:voxel_num]

    def point2voxel_jit(self, points: np.ndarray):
        """
        :param points: 传入的点云数据, shape: (N, num_point_features)
        :param mean: 是否进行平均 (这个参数为看 spcov 里面有, 但是具体到 HEAL 上没有用到, 所以就不写这部分的代码了)
        """
        if points.shape[1] != self.num_point_features:
            Exception(f"{points.shape[1]=} must be equal to {self.num_point_features=}.")
        return point2voxel_jit(
            self.ndim,
            self.dtype,
            self.vsize,
            self.coors_range,
            self.num_point_features,
            self.max_num_voxels,
            self.max_num_points_per_voxel,
            self.grid_size,
            points
        )

测试代码

import os
import time

import loguru
import numpy as np
from cumm import tensorview as tv
from spconv.utils import Point2VoxelCPU3d
from pypcd4 import PointCloud

from src.point2voxel_numpy_3d import Point2VoxelNumpy3D


def read_pcd_files(folder_path):
    """
    递归地读取指定文件夹及其子文件夹中的所有 .pcd 文件。

    Args:
        folder_path (str): 要读取的文件夹路径。

    Returns:
        list: 包含所有找到的 .pcd 文件完整路径的列表。
    """
    pcd_file_paths = []
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(".pcd"):
                pcd_file_path = os.path.join(root, file)
                pcd_file_paths.append(pcd_file_path)
    return pcd_file_paths


def read_pcd(pcd_path):
    pcd = PointCloud.from_path(pcd_path)
    return pcd.numpy()  # 强度信息没有像 HEAL 那样进行处理


if __name__ == "__main__":
    path_list = read_pcd_files(os.path.expanduser("~/Documents/OPV2V"))
    vsize_xyz = np.array([0.4, 0.4, 4])
    coors_range_xyz = np.array([-102.4, -51.2, -3, 102.4, 51.2, 1])
    max_num_points_per_voxel = 32
    num_point_features = 4
    max_num_voxels = 32000

    for pcd_file_path in path_list:
        pcd_np = read_pcd(pcd_file_path)

        spconv_time_start = time.time()
        spconv_voxel_generator = Point2VoxelCPU3d(
            vsize_xyz=vsize_xyz,
            coors_range_xyz=coors_range_xyz,
            max_num_points_per_voxel=max_num_points_per_voxel,
            num_point_features=num_point_features,
            max_num_voxels=max_num_voxels,
        )
        pcd_tv = tv.from_numpy(pcd_np)
        spconv_voxels, spconv_coordinates, spconv_num_points = spconv_voxel_generator.point_to_voxel(pcd_tv)
        spconv_voxels, spconv_coordinates, spconv_num_points = (
            spconv_voxels.numpy(),
            spconv_coordinates.numpy(),
            spconv_num_points.numpy(),
        )

        spconv_time_end = time.time()

        numpy_time_start = time.time()
        numpy_voxel_generator = Point2VoxelNumpy3D(
            vsize_xyz, coors_range_xyz, num_point_features, max_num_voxels, max_num_points_per_voxel
        )
        numpy_voxels, numpy_coordinates, numpy_num_points = numpy_voxel_generator.point2voxel(pcd_np)
        numpy_time_end = time.time()
        loguru.logger.info(
            f"检查生成的 voxels 是否相等: {np.allclose(spconv_voxels, numpy_voxels, rtol=1e-5, atol=1e-8, equal_nan=True)}"
        )
        loguru.logger.info(
            "检查生成的 coordinates 是否相等:"
            f" {np.allclose(spconv_coordinates, numpy_coordinates, rtol=1e-5, atol=1e-8, equal_nan=True)}"
        )
        loguru.logger.info(
            "检查生成的 num_points 是否相等:"
            f" {np.allclose(spconv_num_points, numpy_num_points, rtol=1e-5, atol=1e-8, equal_nan=True)}"
        )
        break

    loguru.logger.info(f"spconv: {spconv_time_end - spconv_time_start} seconds")
    loguru.logger.info(f"numpy: {numpy_time_end - numpy_time_start} seconds")

测试结果

使用 jit
在这里插入图片描述
未使用 jit
在这里插入图片描述
第二种方法
还存在问题, 所以需要进一步改进

其他问题

用 numpy 比较了一下, 发现得到的两种 voxel 并不相等, 个人认为是, 浮点数精度上的问题, 如果调试的时候打印出来, 两者还是相等的. 但为了严谨还需要进一步研究, 确保两者在一定精度范围内确实是相等的.
可以发现, 使用了 jit 的代码要明显块于原始代码, 然而即便如此依然远远落后于 spconv

评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值