姿态估计2-07：PVNet(6D姿态估计)-源码无死角解析（3）-模型总体结构

以下链接是个人关于PVNet(6D姿态估计) 所有见解，如有错误欢迎大家指出，我会第一时间纠正。有兴趣的朋友可以加微信：a944284742相互讨论技术。若是帮助到了你什么，一定要记得点赞！因为这是对我最大的鼓励。
姿态估计2-00：PVNet(6D姿态估计)-目录-史上最新无死角讲解

前言

通过上一篇博客我们已经知道网络是如何获取数据，以及数据对应的标签。该小姐主要对网络模型的总体架构进行讲解。我们在train_net.py文件中可以看到如下代码：

    network = make_network(cfg)

这里就是对网路模型的构建，如果是进行性训练，其会调用到lib/train/trainers/pvnet.py文件中的class NetworkWrapper。该类的注释如下：

import torch.nn as nn
from lib.utils import net_utils
import torch


class NetworkWrapper(nn.Module):
    def __init__(self, net):
        super(NetworkWrapper, self).__init__()

        self.net = net

        self.vote_crit = torch.nn.functional.smooth_l1_loss
        self.seg_crit = nn.CrossEntropyLoss()

    def forward(self, batch):
        # 输入图像像素，获得语义分割的结果(为一个mask)，以及论文中的vectors
        output = self.net(batch['inp'])

        # 记录训练过程中的loss
        scalar_stats = {}
        loss = 0

        # 如果batch['meta']中包含了pose_test
        if 'pose_test' in batch['meta'].keys():
            loss = torch.tensor(0).to(batch['inp'].device)
            return output, loss, {}, {}

        # 使用语义分割出来的mask当作weight
        weight = batch['mask'][:, None].float()

        # 只对mask(目标所在区域)进行vector loos计算
        vote_loss = self.vote_crit(output['vertex'] * weight, batch['vertex'] * weight, reduction='sum')

        vote_loss = vote_loss / weight.sum() / batch['vertex'].size(1)
        scalar_stats.update({'vote_loss': vote_loss})
        loss += vote_loss

        # 计算语义标签mask的loss
        mask = batch['mask'].long()
        seg_loss = self.seg_crit(output['seg'], mask)
        scalar_stats.update({'seg_loss': seg_loss})
        loss += seg_loss

        # 记录loss
        scalar_stats.update({'loss': loss})
        image_stats = {}

        return output, loss, scalar_stats, image_stats

代码简介

其上的：

output = self.net(batch['inp'])

本调试代码使用的是lib/networks/pvnet/resnet18.py，该resnet18为作者修改过后的resnet18，具体修改如下（在论文中有提到）：

1.当网络的feature map的大小为$H=8×W=8$时，我们不再通过丢弃后续的pooling层对feature map进行downsample
2.为了保持接收域不变，用合适的扩张的卷积替换后续的卷积
3.将原ResNet-18中全连接层替换为卷积层

修改过后的网络会有两个输出，一个为语义分割的mask，另外一个为vectors。他们的分辨率和输入图像的大小是一样的，形状分别为：
mask[b,2,H,W] ：这里的b表示batch_size，2表示为一个二分类。
vectors[b,18,H,W]：因为有9个关键点(包含一个中心关键点)，每个像素都要预测其趋向于每个关键点的方向。
 
对于loss的计算，也分为连个部分：
1.语义分割，使用nn.CrossEntropyLoss()，因为该属于分类系列
2.对于vectors只用torch.nn.functional.smooth_l1_loss，在前面的章节，我们以及知道如何去获取vectors对应的标签，预测出来的vectors越接近标签则越好。

大叫要注意的是，对于vectors计算loss的时候，只需要对目标物体存在的像素做反向传播即可，背景是没有作反向传播的，该代码的核心为：

 weight = batch['mask'][:, None].float()

下篇博客开始会为大家讲解一些细节上的，希望大家持续关注。如ransac的投票策略等等。