Loading

『记录』Python项目中的extensions

关于python项目中的extensions

在用mmdetection3d的时候显然许多经典的操作不是写在自己的python里，而是通过叫做ops的东西导入在python里用的，希望用这一篇把这玩意理清楚。

以voxelization操作为例，其链接关系为：

点击查看代码 
# mmdet3d的voxelnet.py中
from mmcv.ops import Voxelization
self.voxel_layer = Voxelization(**voxel_layer)

# mmcv的ops的voxelize.py中
class Voxelization(nn.Module):
...
return voxelization(input, self.voxel_size, self.point_cloud_range,
                            self.max_num_points, max_voxels,
                            self.deterministic)
# 同文件中
class _Voxelization(Function):
...
if max_points == -1 or max_voxels == -1:
	coors = points.new_zeros(size=(points.size(0), 3), dtype=torch.int)
	ext_module.dynamic_voxelize_forward(
	    points,
	    torch.tensor(voxel_size, dtype=torch.float),
	    torch.tensor(coors_range, dtype=torch.float),
	    coors,
	    NDim=3)
		return coors
else:
	voxels = points.new_zeros(size=(max_voxels, max_points, points.size(1)))
  coors = points.new_zeros(size=(max_voxels, 3), dtype=torch.int)
  num_points_per_voxel = points.new_zeros(
      size=(max_voxels, ), dtype=torch.int)
  voxel_num = torch.zeros(size=(), dtype=torch.long)
  ext_module.hard_voxelize_forward(
      points,
      torch.tensor(voxel_size, dtype=torch.float),
      torch.tensor(coors_range, dtype=torch.float),
      voxels,
      coors,
      num_points_per_voxel,
      voxel_num,
      max_points=max_points,
      max_voxels=max_voxels,
      NDim=3,
      deterministic=deterministic)
  # select the valid voxels
  voxels_out = voxels[:voxel_num]
  coors_out = coors[:voxel_num]
  num_points_per_voxel_out = num_points_per_voxel[:voxel_num]
  return voxels_out, coors_out, num_points_per_voxel_out

voxelization = _Voxelization.apply

# 同文件中
from ..utils import ext_loader
ext_module = ext_loader.load_ext(
    '_ext', ['dynamic_voxelize_forward', 'hard_voxelize_forward'])

# mmcv的utils的ext_loader.py中
def load_ext(name, funcs):
        ext = importlib.import_module('mmcv.' + name)
        for fun in funcs:
            assert hasattr(ext, fun), f'{fun} miss in module {name}'
        return ext

所以大概过程是这样的：

  • 在mmcv的文件中，存储了很多非python的.cpp/.h/.cu的文件（以及貌似还有pybind.cpp来帮助），其大概叫做(cuda) operators. python的package安装时，在其setup.py中就使用setup这个东西的一些写法来安装这些“Extensions”（用到了torch.utils.cpp_extension中的CppExtension和CUDAExtension，其是pytorch所写的方便我们返回一个setuptools.Extension（这个setuptools应该也就是目前setup.py中的规范用法））（有点类似CMakeList？）（估计将其编译到了.so之类的library中？）
  • 于是，在python中即可调用这些方法，比如体素化的操作，mmcv通过importlib导入自己的"mmcv._ext"，调用里面的dynamic_voxelize_forward或者hard_voxelize_forward，自己包装了一下自己的extension在这里，成为了一个_Voxelization类，其继承pytorch的Function类关于这个pytorch的Function，其介绍为：“To create a custom autograd.Function, subclass this class and implement the forward and backward static methods. Then, to use your custom op in the forward pass, call the class method apply. Do not call forward directly.”，所以觉得应该是一个如果你要导ops来用的一个规范用法。另外一点是，这个包装根据判断输入max_points和max_voxels这两个在HV中的两个主要信息的值自动判断用HV还是DV，如果都是-1，就用DV（也就是后面我们用的时候想用DV的话就把这两个值置为-1传进来就是了）
  • 然后又用Voxelization这个类继续包装了一下_Voxelization，这个类就是nn.Module类的了。于是，这个就是我们在用的时候可以直接调用的了！往往在写网络的时候，就把**voxel_layer给传进来，那么也就追溯到一开始就是把我们config里写好的keyword argument dict传进来对这个初始化了！在mmdet3d中，voxelization的这一层往往叫做voxel_layer或者如果是像mvx-net这样multi-modality的叫做pts_voxel_layer

上面是自己所猜测的大概思路。真的还有好多细节值得去看，后面看！

posted @ 2022-07-17 22:46  traviscui  阅读(112)  评论(0)    收藏  举报