从代码到BEV：手把手带你复现LSS算法中的关键模块（附PyTorch源码解析）-洪萨配资

从代码到BEV：LSS算法核心模块的PyTorch实现与工程细节剖析

在自动驾驶感知领域，鸟瞰图（BEV）表示正逐渐成为多传感器融合的主流范式。NVIDIA提出的Lift-Splat-Shoot（LSS）算法作为BEV感知的开山之作，其精妙的张量操作和工程实现技巧值得深入探讨。本文将聚焦create_frustum、get_geometry和voxel_pooling三个核心函数，通过可运行的代码示例揭示算法背后的设计哲学。

1. 视锥构建：从2D像素到3D空间的升维艺术

LSS算法的起点是将2D图像特征"抬升"到3D空间。这个过程始于create_frustum函数，它构建了一个参数化的视锥点云。不同于传统点云生成方法，LSS采用了一种内存高效的张量操作策略：

def create_frustum(self): ogfH, ogfW = 128, 352 # 原始图像高度和宽度 fH, fW = 8, 22 # 特征图下采样后的尺寸 # 深度维度构造 (D=41) ds = torch.arange(4, 45, 1).view(-1, 1, 1).expand(-1, fH, fW) # 像素坐标映射 xs = torch.linspace(0, ogfW-1, fW).view(1, 1, fW).expand(-1, fH, fW) ys = torch.linspace(0, ogfH-1, fH).view(1, fH, 1).expand(-1, fH, fW) # 构建3D视锥 (D,H,W,3) frustum = torch.stack((xs, ys, ds), -1) return nn.Parameter(frustum, requires_grad=False)

这个函数的精妙之处在于：

内存视图优化：通过expand操作而非实际复制数据来构建三维坐标
参数化设计：将视锥作为模型的不可训练参数保存
尺度感知：在下采样后的特征图上构建点云，但坐标映射回原始图像空间

提示：实际工程中会使用torch.meshgrid替代部分展开操作，但原始实现保留了更好的内存连续性

2. 坐标系转换：多视角几何的统一表达

get_geometry函数负责将各相机视角的点云转换到统一的自车坐标系。这个过程中涉及多个关键的张量操作：

def get_geometry(self, rots, trans, intrins, post_rots, post_trans): B, N = trans.shape[:2] # batch_size和相机数量 points = self.frustum - post_trans.view(B,N,1,1,1,3) # 逆变换数据增强的旋转 points = torch.inverse(post_rots).view(B,N,1,1,1,3,3) @ points.unsqueeze(-1) # 图像坐标系→相机归一化坐标系 points = torch.cat([ points[...,:2,:] * points[...,2:3,:], points[...,2:3,:] ], dim=-2) # 相机归一化坐标系→自车坐标系 combine = rots @ torch.inverse(intrins) points = combine.view(B,N,1,1,1,3,3) @ points return points.squeeze(-1) + trans.view(B,N,1,1,1,3)

该实现有几个工程亮点：

批量处理：所有相机视角的点云转换在单次前向传播中完成
内存效率：通过view和广播机制避免显式循环
数值稳定：使用torch.inverse而非手动计算逆矩阵

坐标系转换过程中的shape变化轨迹：

操作步骤	张量shape	物理意义
初始frustum	(B,N,D,H,W,3)	图像坐标系下的3D点
数据增强逆变换	(B,N,D,H,W,3,1)	消除数据增强影响
归一化坐标系	(B,N,D,H,W,3,1)	相机归一化坐标
自车坐标系	(B,N,D,H,W,3)	统一的世界坐标

3. 体素池化：高效BEV特征构建的秘诀

voxel_pooling是LSS算法中最具工程挑战的部分，其核心是解决"多对一"投影的高效实现问题。原始实现采用了被称为"cumsum trick"的优化技术：

def voxel_pooling(self, geom_feats, x): B, N, D, H, W, C = x.shape Nprime = B * N * D * H * W # 展平并转换到BEV网格坐标 x = x.reshape(Nprime, C) geom_feats = ((geom_feats - (self.bx - self.dx/2)) / self.dx).long() # 过滤边界外的点 valid = ( (geom_feats[...,0] >= 0) & (geom_feats[...,0] < self.nx[0]) & (geom_feats[...,1] >= 0) & (geom_feats[...,1] < self.nx[1]) & (geom_feats[...,2] >= 0) & (geom_feats[...,2] < self.nx[2]) ) x, geom_feats = x[valid], geom_feats[valid] # 为每个点计算唯一rank ranks = ( geom_feats[...,0] * (self.nx[1] * self.nx[2] * B) + geom_feats[...,1] * (self.nx[2] * B) + geom_feats[...,2] * B + geom_feats[...,3] # batch索引 ) # 排序并应用cumsum trick order = ranks.argsort() x, geom_feats, ranks = x[order], geom_feats[order], ranks[order] x, geom_feats = self.cumsum_trick(x, geom_feats, ranks) # 构建最终BEV特征图 bev_feature = torch.zeros((B, C, *self.nx), device=x.device) bev_feature[geom_feats[:,3], :, geom_feats[:,2], geom_feats[:,0], geom_feats[:,1]] = x return bev_feature.squeeze(2)

cumsum trick的数学原理可以通过一个简单例子理解：

# 假设有5个点的特征和rank值 features = torch.tensor([[1], [2], [3], [4], [5]]) ranks = torch.tensor([0, 1, 2, 2, 3]) # 步骤1：计算累积和 cumsum = features.cumsum(0) # [1,3,6,10,15] # 步骤2：标记rank变化点 keep = torch.ones_like(ranks, dtype=bool) keep[:-1] = ranks[1:] != ranks[:-1] # [True,True,False,True,True] # 步骤3：筛选并差分 filtered = cumsum[keep] # [1,3,10,15] result = torch.cat([filtered[:1], filtered[1:] - filtered[:-1]]) # [1,2,7,5]