【DenseFusion代码详解】网络结构详解

DenseFusion系列代码全讲解目录：【DenseFusion系列目录】代码全讲解+可视化+计算评估指标_Panpanpan！的博客-CSDN博客

这些内容均为个人学习记录，欢迎大家提出错误一起讨论一起学习！

---

代码位置在lib/network.py

这部分是网络的设计。这里有两个主要结构，一个是PoseNet，一个是PoseRefineNet。

PoseNet是主干网络，用于模型初期的训练，模式为train，当模型精度达到一定数值之后（DenseFusion中设置的0.013），就开始进行refine过程，模式改为eval，PoseRefineNet模式改为train。

为了更好地理解，这里从两个主干网络的forward部分进行介绍。

PoseNet

PoseNet网络结构如下：

class PoseNet(nn.Module):
    def __init__(self, num_points, num_obj):
        super(PoseNet, self).__init__()
        self.num_points = num_points
        self.cnn = ModifiedResnet()
        self.feat = PoseNetFeat(num_points)
        
        self.conv1_r = torch.nn.Conv1d(1408, 640, 1)
        self.conv1_t = torch.nn.Conv1d(1408, 640, 1)
        self.conv1_c = torch.nn.Conv1d(1408, 640, 1)

        self.conv2_r = torch.nn.Conv1d(640, 256, 1)
        self.conv2_t = torch.nn.Conv1d(640, 256, 1)
        self.conv2_c = torch.nn.Conv1d(640, 256, 1)

        self.conv3_r = torch.nn.Conv1d(256, 128, 1)
        self.conv3_t = torch.nn.Conv1d(256, 128, 1)
        self.conv3_c = torch.nn.Conv1d(256, 128, 1)

        self.conv4_r = torch.nn.Conv1d(128, num_obj*4, 1) #quaternion
        self.conv4_t = torch.nn.Conv1d(128, num_obj*3, 1) #translation
        self.conv4_c = torch.nn.Conv1d(128, num_obj*1, 1) #confidence

        self.num_obj = num_obj

    def forward(self, img, x, choose, obj):
        out_img = self.cnn(img) #img: torch.Size([bs, 3, 120, 120]) --> out_img：torch.Size([bs, 32, 120, 120])
        
        bs, di, _, _ = out_img.size()

        emb = out_img.view(bs, di, -1) #emb: torch.Size([bs, 32, 14400])
        choose = choose.repeat(1, di, 1) #choose: torch.Size([bs, 1, 500])
        emb = torch.gather(emb, 2, choose).contiguous() #emb: torch.Size([bs, 32, 500])
        
        x = x.transpose(2, 1).contiguous() #points: torch.Size([bs, 500, 3]) -->torch.Size([bs, 3, 500])
        ap_x = self.feat(x, emb) #ap_x: torch.Size([bs, 1408, 500])

        rx = F.relu(self.conv1_r(ap_x)) # torch.Size([bs, 640, 500])
        tx = F.relu(self.conv1_t(ap_x))
        cx = F.relu(self.conv1_c(ap_x))      

        rx = F.relu(self.conv2_r(rx)) # torch.Size([bs, 256, 500])
        tx = F.relu(self.conv2_t(tx))
        cx = F.relu(self.conv2_c(cx))

        rx = F.relu(self.conv3_r(rx)) # torch.Size([bs, 128, 500])
        tx = F.relu(self.conv3_t(tx))
        cx = F.relu(self.conv3_c(cx))

        rx = self.conv4_r(rx).view(bs, self.num_obj, 4, self.num_points) #torch.Size([bs, num_obj, 4, 500])
        tx = self.conv4_t(tx).view(bs, self.num_obj, 3, self.num_points) #torch.Size([bs, num_obj, 3, 500])
        cx = torch.sigmoid(self.conv4_c(cx)).view(bs, self.num_obj, 1, self.num_points) #torch.Size([bs, num_obj, 1, 500])
        
        b = 0
        out_rx = torch.index_select(rx[b], 0, obj[b])
        out_tx = torch.index_select(tx[b], 0, obj[b])
        out_cx = torch.index_select(cx[b], 0, obj[b])
        
        out_rx = out_rx.contiguous().transpose(2, 1).contiguous()
        out_cx = out_cx.contiguous().transpose(2, 1).contiguous()
        out_tx = out_tx.contiguous().transpose(2, 1).contiguous()
        
        return out_rx, out_tx, out_cx, emb.detach()

该部分的输入是预处理之后随机选择的500个点云（代码中的x）、物体的image crop（代码中的img）、随机选取的像素索引（代码中的choose）、物体的类别编号（代码中的obj）。

forward里面可以看到网络的结构。第一行out_img = self.cnn(img)输入是img，也就是RGB图像，大小为[bs,3,h,w]，self.cnn是ModifiedResnet()，为提取颜色特征的网络，网络结构如下：

psp_models = {
    'resnet18': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet18'),
    'resnet34': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet34'),
    'resnet50': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet50'),
    'resnet101': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet101'),
    'resnet152': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet152')
}
class ModifiedResnet(nn.Module):

    def __init__(self, usegpu=True):
        super(ModifiedResnet, self).__init__()

        self.model = psp_models['resnet18'.lower()]()
        self.model = nn.DataParallel(self.model)

    def forward(self, x):
        x = self.model(x)
        return x

提取颜色特征的网络是一种编码-解码结构，编码器是ResNet18，解码器是4个上采样层（PSPNet的金字塔形式），上面定义了不同不同参数的网络组合，在训练的时候可以自己选择。其中，nn.DataParallel函数来用多个GPU来加速训练。对应论文网络结构的CNN部分，对RGB图像进行语义分割后的image crop进行特征提取。也就是下面这部分：

下面一行bs, di, _, _ = out_img.size()，这里out_img是提取之后的颜色特征，也就是上图中的color embeddings，bs是批量大小，di是通道数，我们可以用一个例子来理解一下，加入输入的img大小为[3,120,120]，那么CNN中每一层和输出大小可以通过：

        self.cnn = ModifiedResnet()
        summary(self.cnn,(3, 120, 120))

在后面加一行命令来查看。输出结果如下：

最终输出的通道数为32，也就是通过CNN提取了32维度的颜色特征，那么di就等于32。下面一行

emb = out_img.view(bs, di, -1)，是将输出的特征转换成[bs,32,-1]维度的特征，其中-1用来自动计算剩余维度，比如上面例子输出是[bs,32,120,120]，那么转换之后的emb就是[bs,32,14400]，相当于把高和宽拉长成一维向量了。

下面一行choose = choose.repeat(1, di, 1)，是将choose复制di遍，也就是每个通道都复制一个，这里的choose表示随机选取的点云的index，为数据预处理是返回的点云索引。比如输入的choose大小为[bs,1,500]也就是随机选取的500个点云，然后进行repeat操作后大小为[bs,32,500]。

下面一行emb = torch.gather(emb, 2, choose).contiguous()，首先gather表示的是收集输入的特定维度指定位置的数值，其中dim=2表示在深度维度上。也就是选取choose点云对应位置的颜色特征。当调用contiguous()时，会强制拷贝一份tensor，让它的布局和从头创建的一模一样，但是两个tensor完全没有联系，即改变括号中emb的值，括号外emb的值不会改变。最后输出emb大小为[bs,32,500]。

下面一行x = x.transpose(2, 1).contiguous()，x就是随机选取的500个点云，大小为[bs,500,3]，用transpose交换第1维和第2维，输出[bs,3,500]。代表点云数据。

下面一行ap_x = self.feat(x, emb)，这里self.feat = PoseNetFeat(num_points)，稠密融合就是在这里完成的，PoseNetFeat结构如下：

class PoseNetFeat(nn.Module):
    def __init__(self, num_points):
        super(PoseNetFeat, self).__init__()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)

        self.e_conv1 = torch.nn.Conv1d(32, 64, 1)
        self.e_conv2 = torch.nn.Conv1d(64, 128, 1)

        self.conv5 = torch.nn.Conv1d(256, 512, 1)
        self.conv6 = torch.nn.Conv1d(512, 1024, 1)

        self.ap1 = torch.nn.AvgPool1d(num_points)
        self.num_points = num_points
    def forward(self, x, emb): #x: torch.Size([bs, 3, 500]) emb: torch.Size([bs, 32, 500])
        x = F.relu(self.conv1(x))  #x: torch.Size([bs, 64, 500])
        emb = F.relu(self.e_conv1(emb)) #emb: torch.Size([bs, 64, 500])
        pointfeat_1 = torch.cat((x, emb), dim=1) #pointfeat_1: torch.Size([bs, 64+64=128, 500])

        x = F.relu(self.conv2(x)) #x: torch.Size([bs, 128, 500])
        emb = F.relu(self.e_conv2(emb)) #emb: torch.Size([bs, 128, 500])
        pointfeat_2 = torch.cat((x, emb), dim=1) #pointfeat_2: torch.Size([bs, 128+128=256, 500])

        x = F.relu(self.conv5(pointfeat_2)) #torch.Size([bs, 256, 500]) --> torch.Size([bs, 512, 500]) 
        x = F.relu(self.conv6(x)) #torch.Size([bs, 512, 500]) --> torch.Size([bs, 1024, 500]) 

        ap_x = self.ap1(x) #torch.Size([bs, 1024, 500]) --> torch.Size([bs, 1024, 1]) 

        ap_x = ap_x.view(-1, 1024, 1).repeat(1, 1, self.num_points) #torch.Size([bs, 1024, 1]) --> torch.Size([bs, 1024, 500]) 
        return torch.cat([pointfeat_1, pointfeat_2, ap_x], 1) #128 + 256 + 1024

来看forward函数，输入的x为点云数据，大小为torch.Size([bs, 3, 500]) ，emb为对应颜色特征，大小为torch.Size([bs, 32, 500])。首先对x和emb分别使用1*1卷积和relu激活函数，输出64维度的特征；然后在深度维上将x和emb融合（cat操作）形成pointfeat_1（大小为torch.Size([bs, 128, 500])）；然后继续对x和emb使用1*1卷积和relu激活函数，输出128维度的特征；然后在深度维上进行融合形成pointfeat_2（大小为torch.Size([bs, 256, 500])）；然后对pointfeat_2使用1*1卷积—relu——1*1卷积—relu；输出大小为torch.Size([bs, 1024, 500])的特征，然后使用全局平均池化（self.ap1 = torch.nn.AvgPool1d(num_points)），每num_points列求平均，这里num_points就等于500，也就是每个通道都变成了平均值，输出ap_x大小为torch.Size([bs, 1024, 1])，也就是论文中的global feature全局特征；然后复制500份，变成torch.Size([bs, 1024, 500])；最后，将pointfeat_1，pointfeat_2，ap_x在通道维上融合，输出torch.Size([bs, 128+256+1024, 500])。

这里PoseNetFeat的过程对应论文里面以下部分：

这里体现了PointNet的思想，图上面PointNet画在点云数据后面，但实际上是通过这两部分体现的：

 图上color就是代码里面的emb，point就是代码里面的x。

再回到PoseNet，接下来的一行是rx = F.relu(self.conv1_r(ap_x))，从这一行到以下这部分代码：

是为每个像素回归r、t、c，r代表旋转，t代表平移，c代表置信度。输入ap_x大小为torch.Size([bs, 1408, 500])，连续用4个1*1卷积最后得出rx大小为torch.Size([bs, num_obj*4, 500])，tx大小为torch.Size([bs, num_obj*3, 500])，cx大小为torch.Size([bs, num_obj*1, 500])，其中num_obj为物体类别个数，linemod数据集为13。然后对rx和tx进行重构得到大小分别为torch.Size([bs, num_obj, 4, 500])和torch.Size([bs, num_obj, 3, 500])，rx和tx构成估计的姿态，每一个像素都有每个类别对应的姿态，cx重构成torch.Size([bs, num_obj, 1, 500])用sigmoid激活作为置信度，每一像素都有一个置信度。

下面这一部分：

obj是输入的物体类别，大小为torch.Size([bs, 1])，比如物体是第2个类别，那么obj=tensor([[2]])，obj[0] = tensor([2])，index_select函数选取指定的tensor，输入为rx[0] ，也就是不考虑批量（因为这里批量都是1），那么0维度就是num_obj维度，选取obj[0]类别对应的数值。通俗来说，每个像素为每个类别都预测了pose，这里就是要找到对应类别的pose，然后交换后两个维度（为了和输入对应）输出。这里没有使用投票法，而是输出所有像素预测的pose，投票的步骤在后续算loss的时候才进行。

该部分的输出为预测的每像素的旋转r、平移t、置信度c、随机选择之后的500个像素的RGB图像。

以上就是整个PoseNet的内容，下面是PoseRefineNet的部分。

PoseRefineNet

代码如下：

class PoseRefineNet(nn.Module):
    def __init__(self, num_points, num_obj):
        super(PoseRefineNet, self).__init__()
        self.num_points = num_points
        self.feat = PoseRefineNetFeat(num_points)
        
        self.conv1_r = torch.nn.Linear(1024, 512)
        self.conv1_t = torch.nn.Linear(1024, 512)

        self.conv2_r = torch.nn.Linear(512, 128)
        self.conv2_t = torch.nn.Linear(512, 128)

        self.conv3_r = torch.nn.Linear(128, num_obj*4) #quaternion
        self.conv3_t = torch.nn.Linear(128, num_obj*3) #translation

        self.num_obj = num_obj

    def forward(self, x, emb, obj):
        bs = x.size()[0]
        
        x = x.transpose(2, 1).contiguous()
        ap_x = self.feat(x, emb)

        rx = F.relu(self.conv1_r(ap_x))
        tx = F.relu(self.conv1_t(ap_x))   

        rx = F.relu(self.conv2_r(rx))
        tx = F.relu(self.conv2_t(tx))

        rx = self.conv3_r(rx).view(bs, self.num_obj, 4)
        tx = self.conv3_t(tx).view(bs, self.num_obj, 3)

        b = 0
        out_rx = torch.index_select(rx[b], 0, obj[b])
        out_tx = torch.index_select(tx[b], 0, obj[b])

        return out_rx, out_tx

该部分的输入为由上一步预测的[R|t]（经过loss计算选取的置信度最大的姿态）转换之后的点云（代码中的x）和上一步PoseNet输出的500像素的RGB图像。

forward中也是一样的思路，但这里的emb以及是经过choose之后的了，也没有使用PSPNet提取颜色特征，直接用self.feat = PoseRefineNetFeat(num_points)进行融合，PoseRefineNetFeat结构如下：

class PoseRefineNetFeat(nn.Module):
    def __init__(self, num_points):
        super(PoseRefineNetFeat, self).__init__()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)

        self.e_conv1 = torch.nn.Conv1d(32, 64, 1)
        self.e_conv2 = torch.nn.Conv1d(64, 128, 1)

        self.conv5 = torch.nn.Conv1d(384, 512, 1)
        self.conv6 = torch.nn.Conv1d(512, 1024, 1)

        self.ap1 = torch.nn.AvgPool1d(num_points)
        self.num_points = num_points

    def forward(self, x, emb):
        x = F.relu(self.conv1(x))
        emb = F.relu(self.e_conv1(emb))
        pointfeat_1 = torch.cat([x, emb], dim=1)

        x = F.relu(self.conv2(x))
        emb = F.relu(self.e_conv2(emb))
        pointfeat_2 = torch.cat([x, emb], dim=1)

        pointfeat_3 = torch.cat([pointfeat_1, pointfeat_2], dim=1)

        x = F.relu(self.conv5(pointfeat_3))
        x = F.relu(self.conv6(x))

        ap_x = self.ap1(x)

        ap_x = ap_x.view(-1, 1024)
        return ap_x

输入点云x和rgb图像emb，首先分别用1*1卷积和relu激活函数，输出64维度的特征；然后在深度维上将x和emb融合（cat操作）形成pointfeat_1（大小为torch.Size([bs, 128, 500])）；然后继续对x和emb使用1*1卷积和relu激活函数，输出128维度的特征；然后在深度维上进行融合形成pointfeat_2（大小为torch.Size([bs, 256, 500])）；

以上都与之前PoseNet部分相同。

下面把pointfeat_1和pointfeat_2进行cat，形成pointfeat_3（大小为torch.Size([bs,128+256, 500])），然后对pointfeat_3使用1*1卷积—relu—1*1卷积—relu；输出大小为torch.Size([bs, 1024, 500])的特征，然后使用全局平均池化输出ap_x大小为torch.Size([bs, 1024, 1])，之后没有复制500份，而是转换成大小为torch.Size([bs, 1024])之后直接输出。

回到PoseRefineNet部分，得到全局特征之后，用该特征回归旋转r、平移t和置信度c，但这里只有唯一的像素，因为refine过程没有必要再对每个像素投票，它是一个姿态矫正的过程。

该部分的输出为优化之后的旋转R和平移t。

总结

该部分是网络的设计部分。包括PoseNet主干网络部分和PoseRefineNet后续的迭代自优化部分。DenseFusion中分别用两个主干网络提取颜色特征和几何特征，为什么不用一个网络？因为这两种数据来自不同的数据空间，并且存在一些投影分解问题（近年有研究发现的）。对于颜色特征，采用编码解码结构，对于几何特征，实际用到的是1*1卷积，PointNet的局部+全局融合的思想在融合部分体现。融合部分采用像素级稠密融合方式，进行了两次融合。第一次形成局部特征，然后用局部特征提取全局特征，再把全局特征融合到局部特征中。最后用这种稠密特征回归姿态和置信度。