BEV感知经典LSS(Lift-Splat-Shoot)论文注释代码，三十分钟全搞定

MinyounZhang

已于 2023-11-26 22:38:45 修改

阅读量1.7k

点赞数 35

文章标签：深度学习自动驾驶计算机视觉神经网络

于 2023-11-24 11:27:25 首次发布

本文链接：https://blog.csdn.net/MinyounZhang/article/details/134582017

版权

家人们大家好，我是老张。今天分享一波自动驾驶BEV（Bird-eye view，鸟瞰图视角）感知范式的开创性工作，NVIDIA提出的LSS（Lift-Splat-Shoot）。尽管时至今日，构建BEV特征的方式逐渐丰富；显式构建BEV特征也不再是必须，但阅读和思考本文仍然会有很多收获。

我的分享包括代码和中文注释，没有太多文字描述与公式，争取三十分钟人人都能看明白，现在发车。论文和官方Repo链接：

arxiv.org/pdf/2008.05711.pdf nv-tlabs/lift-splat-shoot: Lift, Splat,
https://github.com/nv-tlabs/lift-splat-shoot

LSS提出了一种将多视角的相机图像融合在BEV空间下的编码方法，其中：

Lift是通过预测深度信息，将2D图像编码到3D空间；
Splat是将3D特征的高度拍扁为BEV特征；
Shoot是运动规划。

在本工作中，BEV特征空间是以自车为中心的栅格，X轴Y轴最远检测距离各50米，Z轴各10米。这个100 * 100 * 20的空间就是BEV特征需要表征的，栅格内每个位置是空间中该位置的特征。

xbound=[-50.0, 50.0, 0.5],   # x方向网格
ybound=[-50.0, 50.0, 0.5],   # y方向网格
zbound=[-10.0, 10.0, 20.0],  # z方向网格
dbound=[4.0, 45.0, 1.0],     # 深度方向网格

将2D图像编码到BEV视角下，需要模型有优质的深度信息，距离预测的准物体在BEV特征中的位置才会准。关于如何把环视摄像机的图片编码到BEV空间下，LSS采用最直接的方式：显式的预测深度信息；另一范式如BEVFormer等则使用Transformer结构，隐式的通过BEV特征空间的query查询2D图像中的特征。
在这里插入图片描述

来看模型流水线图，其实Extrinsics+Splat，完整流水线可以分为5部分：

1.生成视锥点云，代表特征图中元素的相机坐标
2.根据相机内外参数、数据增强参数，将视锥点云从相机坐标映射到世界坐标
3.CamEncode网络提取2D图片的深度特征矩阵
4.⭐深度特征绑定视锥点云，Voxel Pooling为BEV特征
5.BEV特征编码，输出

一、2D图片 - > 深度特征

我们先从最容易理解的【3】开始说起，不考虑batch和摄像头数量：一张2D图片的张量维度为【3, W, H】；经过特征提取网络，输出张量为【41+64, Wf, Hf】，Wf和Hf是下采样后特征图的高宽。

输出通道数量105是个特殊值，这里前41个通道表示深度取值范围从4米到45米，后64的通道存放当前点的特征。通道拆分做内积得到64 * 41的矩阵，表示该特征点分别在4 ~ 45米每个深度下的特征值。所以模型最终输出张量维度为【 64, 41, Wf, Hf】，对应代码20 ~21行

下方是论文给的图，把矩阵放到图左边或许会更直观些；本段代码实现特征提取+输出深度特征矩阵，不想看也没关系
在这里插入图片描述

class CamEncode(nn.Module):  # 提取图像特征，进行图像深度编码
    def __init__(self, D, C, downsample):
        super(CamEncode, self).__init__()
        self.D = D  # 41 深度区间【4-45】
        self.C = C  # 64 点的特征向量维度

        self.trunk = EfficientNet.from_pretrained("efficientnet-b0")  # efficientnet 提取特征

        self.up1 = Up(320+112, 512)  # 上采样模块，输入320+112（多尺度融合），输出通道512
        self.depthnet = nn.Conv2d(512, self.D + self.C, kernel_size=1, padding=0)  # 1x1卷积调整通道数

    def get_depth_dist(self, x, eps=1e-20):  # 深度维计算softmax，得到每个像素不同深度的概率
        return x.softmax(dim=1)

    def get_depth_feat(self, x): 
        x = self.get_eff_depth(x)  # 使用efficientnet提取特征  x: BN x 512 x 8 x 22
    
        x = self.depthnet(x)  # 1x1卷积变换维度  x: BN x 105(C+D) x 8 x 22
                                                                                    
        depth = self.get_depth_dist(x[:, :self.D])  # 第二个维度的前D个作为深度维，进行softmax  depth: BN x 41 x 8 x 22
        new_x = depth.unsqueeze(1) * x[:, self.D:(self.D + self.C)].unsqueeze(2)  # 将深度概率分布和特征通道利用广播机制相乘  
        
        return depth, new_x #  new_x: BN x 64 x 41 x 8 x 22

    def get_eff_depth(self, x):  # 使用efficientnet提取特征
        # adapted from https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py#L231
        endpoints = dict()

        # Stem
        x = self.trunk._swish(self.trunk._bn0(self.trunk._conv_stem(x)))  #  x: BN x 32 x 64 x 176
        prev_x = x 

        # Blocks
        for idx, block in enumerate(self.trunk._blocks):
            drop_connect_rate = self.trunk._global_params.drop_connect_rate
            if drop_connect_rate:
                drop_connect_rate *= float(idx) / len(self.trunk._blocks) # scale drop connect_rate
            x = block(x, drop_connect_rate=drop_connect_rate)
            if prev_x.size(2) > x.size(2):
                endpoints['reduction_{}'.format(len(endpoints)+1)] = prev_x
            prev_x = x

        # Head
        endpoints['reduction_{}'.format(len(endpoints)+1)] = x  # x: BN x 320 x 4 x 11
        x = self.up1(endpoints['reduction_5'], endpoints['reduction_4'])  # 对endpoints[4]上采样，然后和endpoints[5] concat 在一起
        return x  # x: 24 x 512 x 8 x 22

    def forward(self, x):
        depth, x = self.get_depth_feat(x)  # depth: B*N x D x fH x fW(24 x 41 x 8 x 22)  x: B*N x C x D x fH x fW(24 x 64 x 41 x 8 x 22)

        return x

二、相机坐标->世界坐标的映射关系

然后要回到【1】生成视锥点云。我们得到了单张2D图片的深度特征矩阵。这些特征如何映射到世界坐标系下，和BEV视角的栅格关联上呢？第一步，先建立特征图在世界坐标系下坐标位置的视锥点云，也就是深度特征每个点的在相机坐标系的位置。
CamEncode网络输出维度是【 64, 41, fW, fH】，表示高宽深分别为【fH，fW，41】的立方体中，每个位置特征64维；视锥点云维度【41 x fH x fW x 3】中，3代表XYZ三条坐标轴，因此视锥点云存储立方体每个位置在相机坐标系的XYZ坐标。有些类似位置编码对吧。

def create_frustum(self):
    # make grid in image plane
    ogfH, ogfW = self.data_aug_conf['final_dim']
    fH, fW = ogfH // self.downsample, ogfW // self.downsample # fw和fh是特征图长宽
    ds = torch.arange(*self.grid_conf['dbound'], dtype=torch.float).view(-1, 1, 1).expand(-1, fH, fW)
    D, _, _ = ds.shape # D是深度区间41，【4-45】
    #【0, fW, 2fW, 3fW ..., ogfW-2fW, ogfW-fW, ogfW】
    xs = torch.linspace(0, ogfW - 1, fW, dtype=torch.float).view(1, 1, fW).expand(D, fH, fW)
    ys = torch.linspace(0, ogfH - 1, fH, dtype=torch.float).view(1, fH, 1).expand(D, fH, fW)

    # D x fH x fW x 3
    frustum = torch.stack((xs, ys, ds), -1)
    return nn.Parameter(frustum, requires_grad=False) # 不计算梯度

步骤【2】坐标系转换。这些特征点的在相机坐标系的位置，需要通过相机内外参数映射到世界坐标系下。代码中的映射涉及仿射变化和相机参数，如果没接触过，跳过这部分也不会影响理解论文。注意有6张不同方向的环视摄像头图片，所以内外参是6部相机的内外参数组成。

  def get_geometry(self, rots, trans, intrins, post_rots, post_trans):
        """    
            rots：相机外参旋转, trans：相机外参平移, intrins：相机内参, post_rots：数据增强旋转, post_trans：数据增强平移
            Determine the (x,y,z) locations (in the ego frame) of the points in the point cloud.
            Returns B x N x D x H/downsample x W/downsample x 3
        """
        B, N, _ = trans.shape

        # undo post-transformation
        # B x N x D x H x W x 3，先还原数据增强中旋转和平移对坐标的影响
        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).matmul(points.unsqueeze(-1))

        # cam_to_ego，利用相机内外参
        points = torch.cat((points[:, :, :, :, :, :2] * points[:, :, :, :, :, 2:3],
                            points[:, :, :, :, :, 2:3]
                            ), 5)
        combine = rots.matmul(torch.inverse(intrins)) 
        points = combine.view(B, N, 1, 1, 1, 3, 3).matmul(points).squeeze(-1)
        points += trans.view(B, N, 1, 1, 1, 3)

        return points #维度不变，坐标值从相机坐标系->世界坐标系

三、BEV池化

至此，我们将6个环视摄像头拍摄的图片，通过显式的计算41个深度下的特征，将这【6，41 ，fW， fH】个特征点一一放置在了早先准备好的BEV特征空间下。终于可以到达【4】，构建BEV特征。

在这张像棋盘一样的BEV特征空间中，有位置放了多个特征点，有的位置没放特征点，因此需要【4】在每个栅格进行BEV池化，得到尺寸统一的BEV特征。

官方代码的实现是：

先根据特征点的XYZ坐标和batch，计算每个点的索引值；索引值相同的点位于同一个栅格中；【代码26-30】
根据索引值排序，则索引值变为有序数组，形如【1.2，1.4，2，4，4，5 … 】；
只需“遍历”索引值，将相同索引值的位置求和，完成池化，形如【1.2+1.4，2，4+4，5 … 】

下方voxel_pooling函数准备索引，BEV池化在QuickCumsum类中进行；

def voxel_pooling(self, geom_feats, x):
        # geom_feats: B x N x D x H x W x 3 (4 x 6 x 41 x 8 x 22 x 3)
        # x: B x N x D x fH x fW x C(4 x 6 x 41 x 8 x 22 x 64)

        B, N, D, H, W, C = x.shape  # B: 4  N: 6  D: 41  H: 8  W: 22  C: 64
        Nprime = B*N*D*H*W  # Nprime: 173184

        # flatten x
        x = x.reshape(Nprime, C)  # 将图像展平，一共有 B*N*D*H*W 个点

        # flatten indices
        geom_feats = ((geom_feats - (self.bx - self.dx/2.)) / self.dx).long()  # 将[-50,50] [-10 10]的范围平移到[0,100] [0,20]，计算栅格坐标并取整
        geom_feats = geom_feats.view(Nprime, 3)  # 将像素映射关系同样展平  geom_feats: B*N*D*H*W x 3 (173184 x 3)
        batch_ix = torch.cat([torch.full([Nprime//B, 1], ix,
                             device=x.device, dtype=torch.long) for ix in range(B)])  # 每个点对应于哪个batch
        geom_feats = torch.cat((geom_feats, batch_ix), 1)  # geom_feats: B*N*D*H*W x 4(173184 x 4), geom_feats[:,3]表示batch_id

        # filter out points that are outside box
        # 过滤掉在边界线之外的点 x:0~199  y: 0~199  z: 0
        kept = (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 = x[kept]  # x: 168648 x 64
        geom_feats = geom_feats[kept]

        # get tensors from the same voxel next to each other
        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]  # 给每一个点一个rank值，rank相等的点在同一个batch，并且在在同一个格子里面
        sorts = ranks.argsort()
        x, geom_feats, ranks = x[sorts], geom_feats[sorts], ranks[sorts]  # 按照rank排序，这样rank相近的点就在一起了
        # x: 168648 x 64  geom_feats: 168648 x 4  ranks: 168648

        # cumsum trick
        if not self.use_quickcumsum:
            x, geom_feats = cumsum_trick(x, geom_feats, ranks)
        else:
            x, geom_feats = QuickCumsum.apply(x, geom_feats, ranks)  # 一个batch的一个格子里只留一个点 x: 29072 x 64  geom_feats: 29072 x 4

        # griddify (B x C x Z x X x Y)
        final = torch.zeros((B, C, self.nx[2], self.nx[0], self.nx[1]), device=x.device)  # final: 4 x 64 x 1 x 200 x 200
        final[geom_feats[:, 3], :, geom_feats[:, 2], geom_feats[:, 0], geom_feats[:, 1]] = x  # 将x按照栅格坐标放到final中

        # collapse Z
        final = torch.cat(final.unbind(dim=2), 1)  # 消除掉z维

        return final  # final: 4 x 64 x 200 x 200

torch.autograd.Function是自定义的算子，路径已经变了，能在python层面自定义算子的前向和后向过程。代码中前向过程使用前缀和区间求和；向量错位比较前后元素是否一致，这些计算步骤都使计算过程向量化；前向过程存储了保留下特征的位置，反向过程再使用前缀和技巧，将特征保留位置的梯度分给求和的位置

class QuickCumsum(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, geom_feats, ranks):
        # x: 168648 x 64  geom_feats: 168648 x 4  ranks: 168648 x
        x = x.cumsum(0) # 求前缀和  x: 168648 x 64
        kept = torch.ones(x.shape[0], device=x.device, dtype=torch.bool)  # kept: 168648 x
        kept[:-1] = (ranks[1:] != ranks[:-1])  # rank错位比较，rank[0]!=rank[1],则留下rank[1]

        x, geom_feats = x[kept], geom_feats[kept]  # rank值相等的点只留下最后一个，即一个batch中的一个格子里只留最后一个点 x: 29072  geom_feats: 29072 x 4
        x = torch.cat((x[:1], x[1:] - x[:-1]))  # x错位相减，还原前缀和之前的x，此时点的feature是rank相同点之和，相当于把同一个格子的点特征进行了sum

        # save kept for backward
        ctx.save_for_backward(kept)
        # no gradient for geom_feats
        ctx.mark_non_differentiable(geom_feats)

        return x, geom_feats

    @staticmethod
    def backward(ctx, gradx, gradgeom):
        kept, = ctx.saved_tensors
        back = torch.cumsum(kept, 0)
        back[kept] -= 1
        val = gradx[back]
        return val, None, None

以上是本文主要的创新点。当然鉴于完整性，也奉上BEV特征编码部分代码，总而言之用卷积堆了些计算量，并做了多尺度融合。在这个位置堆计算量是很必要的，因为在此之前所有特征的计算是相机之间独立的，现在所有特征都统一到BEV视角，当然要利用好相机之间特征的语义信息，把BEV特征做大做强。

四、BEV空间下特征再编码

class BevEncode(nn.Module):
    def __init__(self, inC, outC):  # inC: 64  outC: 1
        super(BevEncode, self).__init__()
        # 使用resnet的前3个stage作为backbone
        trunk = resnet18(pretrained=False, zero_init_residual=True)
        self.conv1 = nn.Conv2d(inC, 64, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = trunk.bn1
        self.relu = trunk.relu

        self.layer1 = trunk.layer1
        self.layer2 = trunk.layer2
        self.layer3 = trunk.layer3

        self.up1 = Up(64+256, 256, scale_factor=4)
        self.up2 = nn.Sequential(  # 2倍上采样->3x3卷积->1x1卷积
            nn.Upsample(scale_factor=2, mode='bilinear',
                              align_corners=True),
            nn.Conv2d(256, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, outC, kernel_size=1, padding=0),
        )

    def forward(self, x):  # x: 4 x 64 x 200 x 200
        x = self.conv1(x)  # x: 4 x 64 x 100 x 100
        x = self.bn1(x)
        x = self.relu(x)
        x1 = self.layer1(x)  # x1: 4 x 64 x 100 x 100
        x = self.layer2(x1)  # x: 4 x 128 x 50 x 50
        x = self.layer3(x)  # x: 4 x 256 x 25 x 25
        x = self.up1(x, x1)  # 给x进行4倍上采样然后和x1 concat 在一起  x: 4 x 256 x 100 x 100
        x = self.up2(x)  # 2倍上采样->3x3卷积->1x1卷积  x: 4 x 1 x 200 x 200

        return x