densefusion代码

转载于
https://blog.csdn.net/weixin_43013761/article/details/103053585
代码详细注解
从之前的博客，我相信大家都已经知道，训练代码为tools/train.py，下面时对该代码的详细注解（这里只要随便看看就好，最后面还有总结）

1. tools/train

# --------------------------------------------------------
# DenseFusion 6D Object Pose Estimation by Iterative Dense Fusion
# Licensed under The MIT License [see LICENSE for details]
# Written by Chen
# --------------------------------------------------------

import _init_paths
import argparse
import os
import random
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from datasets.ycb.dataset import PoseDataset as PoseDataset_ycb
from datasets.warehouse.dataset import PoseDataset as PoseDataset_warehouse
from datasets.linemod.dataset import PoseDataset as PoseDataset_linemod
from lib.network import PoseNet, PoseRefineNet
from lib.loss import Loss
from lib.loss_refiner import Loss_refine
from lib.utils import setup_logger
from torchsummary import summary


parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default = 'ycb', help='ycb or warehouse or linemod')
parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir (''YCB_Video_Dataset'' or ''Warehouse_Dataset'' or ''Linemod_preprocessed'')')
parser.add_argument('--batch_size', type=int, default = 8, help='batch size')

# 加载数据的线程数目
parser.add_argument('--workers', type=int, default = 10, help='number of data loading workers')

# 初始学习率
parser.add_argument('--lr', default=0.0001, help='learning rate')

parser.add_argument('--lr_rate', default=0.3, help='learning rate decay rate')

# 初始权重
parser.add_argument('--w', default=0.015, help='learning rate')
# 权重衰减率
parser.add_argument('--w_rate', default=0.3, help='learning rate decay rate')

#
parser.add_argument('--decay_margin', default=0.016, help='margin to decay lr & w')

# 大概是loss到了这个设定的值，则会进行refine模型的训练
parser.add_argument('--refine_margin', default=0.013, help='margin to start the training of iterative refinement')

# 给训练数据添加噪声相关的参数，可以理解为数据增强
parser.add_argument('--noise_trans', default=0.03, help='range of the random noise of translation added to the training data')

# 训练refinenet的时候是连续迭代几次
parser.add_argument('--iteration', type=int, default = 2, help='number of refinement iterations')

# 训练到多少个epoch则停止
parser.add_argument('--nepoch', type=int, default=500, help='max number of epochs to train')

# 是否继续训练posenet模型，继续训练则加载posenet预训练模型
parser.add_argument('--resume_posenet', type=str, default = '',  help='resume PoseNet model')
# 是否继续训练refinenet模型，继续训练则加载refinenet预训练模型
parser.add_argument('--resume_refinenet', type=str, default = '',  help='resume PoseRefineNet model')

parser.add_argument('--start_epoch', type=int, default = 1, help='which epoch to start')
opt = parser.parse_args()


def main():
    opt.manualSeed = random.randint(1, 100)

    # 为CPU随机生成数设定的种子
    random.seed(opt.manualSeed)
    # 为GPU随机生成数设定的种子
    torch.manual_seed(opt.manualSeed)

    # 根据数据集的不同，分别配置其
    # 训练数据的物体种类数目，输入点云的数目，训练模型保存的目录，log保存的目录，起始的epoch数目
    if opt.dataset == 'ycb':
        opt.num_objects = 21 #number of object classes in the dataset
        opt.num_points = 1000 #number of points on the input pointcloud
        opt.outf = 'trained_models/ycb' #folder to save trained models
        opt.log_dir = 'experiments/logs/ycb' #folder to save logs
        opt.repeat_epoch = 1 #number of repeat times for one epoch training
    elif opt.dataset == 'warehouse':
        opt.num_objects = 13
        opt.num_points = 1000
        opt.outf = 'trained_models/warehouse'
        opt.log_dir = 'experiments/logs/warehouse'
        opt.repeat_epoch = 1
    elif opt.dataset == 'linemod':
        opt.num_objects = 13
        opt.num_points = 500
        opt.outf = 'trained_models/linemod'
        opt.log_dir = 'experiments/logs/linemod'
        opt.repeat_epoch = 20
    else:
        print('Unknown dataset')
        return

    # 该处为网络的构建，构建完成之后，能对物体的6D姿态进行预测
    estimator = PoseNet(num_points = opt.num_points, num_obj = opt.num_objects)
    estimator.cuda()
    #summary(estimator,[(3, 120, 120),(500,3),(1,500),(1,)])
    # 对初步预测的姿态进行提炼
    refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_objects)
    refiner.cuda()

    # 对posenet以及refinenet模型的加载，然后标记对应的网络是否已经开始训练过了，以及是否进行衰减
    if opt.resume_posenet != '':
        estimator.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet)))
    if opt.resume_refinenet != '':
        refiner.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet)))
        opt.refine_start = True
        opt.decay_start = True
        opt.lr *= opt.lr_rate
        opt.w *= opt.w_rate
        opt.batch_size = int(opt.batch_size / opt.iteration)
        optimizer = optim.Adam(refiner.parameters(), lr=opt.lr)
    else:
        opt.refine_start = False
        opt.decay_start = False
        optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)

    # 加载对应的训练和验证数据集
    if opt.dataset == 'ycb':
        dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
    elif opt.dataset == 'warehouse':
        dataset = PoseDataset_warehouse('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
    elif opt.dataset == 'linemod':
        dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
    if opt.dataset == 'ycb':
        test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
    elif opt.dataset == 'warehouse':
        test_dataset = PoseDataset_warehouse('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
    elif opt.dataset == 'linemod':
        test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
    testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers)


    opt.sym_list = dataset.get_sym_list()
    #print(opt.sym_list)
    opt.num_points_mesh = dataset.get_num_points_mesh()


    print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<\nlength of the training set: {0}\nlength of the testing set: {1}\nnumber of sample points on mesh: {2}\nsymmetry object list: {3}'.format(len(dataset), len(test_dataset), opt.num_points_mesh, opt.sym_list))

    # loss计算
    criterion = Loss(opt.num_points_mesh, opt.sym_list)
    criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)

    # 初始设置最好模型的loss为无限大
    best_test = np.Inf

    if opt.start_epoch == 1:
        for log in os.listdir(opt.log_dir):
            os.remove(os.path.join(opt.log_dir, log))
    st_time = time.time()


    # 开始循环迭代
    for epoch in range(opt.start_epoch, opt.nepoch):
        # 保存开始开始迭代的log信息
        logger = setup_logger('epoch%d' % epoch, os.path.join(opt.log_dir, 'epoch_%d_log.txt' % epoch))
        logger.info('Train time {0}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + ', ' + 'Training started'))
        train_count = 0
        train_dis_avg = 0.0
        # 判断是否开始训练refine模型
        if opt.refine_start:
            estimator.eval()
            refiner.train()
        else:
            estimator.train()

        optimizer.zero_grad()

        # 每次 epoch 重复训练的次数
        for rep in range(opt.repeat_epoch):
            for i, data in enumerate(dataloader, 0):
                points, choose, img, target, model_points, idx = data
                # points：由深度图计算出来的点云，该点云数据以摄像头主轴参考坐标
                # choose：所选择点云的索引，[bs, 1, 500]
                # img：通过box剪切下来的RGB图像
                # target：根据model_points点云信息，以及旋转偏移矩阵转换过的点云信息[bs,500,3]
                # model_points：目标初始帧（模型）对应的点云信息[bs,500,3]
                # idx：训练图片样本的下标
                points, choose, img, target, model_points, idx = Variable(points).cuda(), \
                                                                 Variable(choose).cuda(), \
                                                                 Variable(img).cuda(), \
                                                                 Variable(target).cuda(), \
                                                                 Variable(model_points).cuda(), \
                                                                 Variable(idx).cuda()

                # 进行预测获得，获得预测的姿态，姿态预测之前的特征向量
                # pred_r: 预测的旋转参数[bs, 500, 4]
                # pred_t: 预测的偏移参数[bs, 500, 3]
                # pred_c: 预测的置信度[bs, 500, 1]，置信度
                #
                pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)

                # 对结果进行评估，计算loss
                loss, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start)

                # 如果已经对refiner模型进行了训练，则进行姿态的提炼预测，对结果进行评估计算dis，并且对dis反向传播
                if opt.refine_start:
                    for ite in range(0, opt.iteration):
                        pred_r, pred_t = refiner(new_points, emb, idx)
                        dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_target, model_points, idx, new_points)
                        dis.backward()
                else:
                    loss.backward()
                train_dis_avg += dis.item()
                train_count += 1

                # log信息存储
                if train_count % opt.batch_size == 0:
                    logger.info('Train time {0} Epoch {1} Batch {2} Frame {3} Avg_dis:{4}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, int(train_count / opt.batch_size), train_count, train_dis_avg / opt.batch_size))
                    optimizer.step()
                    optimizer.zero_grad()
                    train_dis_avg = 0

                # 模型保存
                if train_count != 0 and train_count % 1000 == 0:
                    if opt.refine_start:
                        torch.save(refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format(opt.outf))
                    else:
                        torch.save(estimator.state_dict(), '{0}/pose_model_current.pth'.format(opt.outf))

        print('>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(epoch))


        logger = setup_logger('epoch%d_test' % epoch, os.path.join(opt.log_dir, 'epoch_%d_test_log.txt' % epoch))
        logger.info('Test time {0}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)) + ', ' + 'Testing started'))
        test_dis = 0.0
        test_count = 0

        # 验证模型构建
        estimator.eval()
        refiner.eval()

        for j, data in enumerate(testdataloader, 0):
            # 获得验证模型的的输入数据，并且分配在GPU上运行
            points, choose, img, target, model_points, idx = data
            points, choose, img, target, model_points, idx = Variable(points).cuda(), \
                                                             Variable(choose).cuda(), \
                                                             Variable(img).cuda(), \
                                                             Variable(target).cuda(), \
                                                             Variable(model_points).cuda(), \
                                                             Variable(idx).cuda()
            pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)
            # 对结果进行评估
            _, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start)

            # 如果refine模型已经开始训练，则对该模型也进行评估
            if opt.refine_start:
                for ite in range(0, opt.iteration):
                    pred_r, pred_t = refiner(new_points, emb, idx)
                    dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_target, model_points, idx, new_points)

            test_dis += dis.item()

            # 保存评估的log信息
            logger.info('Test time {0} Test Frame No.{1} dis:{2}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), test_count, dis))

            test_count += 1

        # 计算测试数据的平均dis
        test_dis = test_dis / test_count
        logger.info('Test time {0} Epoch {1} TEST FINISH Avg dis: {2}'.format(time.strftime("%Hh %Mm %Ss", time.gmtime(time.time() - st_time)), epoch, test_dis))
        # 如果该次的测试结果，比之前最好的模型还要好，则保存目前的模型为最好的模型
        if test_dis <= best_test:
            best_test = test_dis
            if opt.refine_start:
                torch.save(refiner.state_dict(), '{0}/pose_refine_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis))
            else:
                torch.save(estimator.state_dict(), '{0}/pose_model_{1}_{2}.pth'.format(opt.outf, epoch, test_dis))
            print(epoch, '>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<')

        # 判断模型测试的结果是否达到，学习率和权重衰减的衰减要求，达到了则进行权重和学习率的衰减
        if best_test < opt.decay_margin and not opt.decay_start:
            opt.decay_start = True
            opt.lr *= opt.lr_rate
            opt.w *= opt.w_rate
            optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)

        # 如果模型没有达到refine_margin的基准（也就是loss比设定的要低），并且refine_start=False，则设定opt.refine_start = True
        # 设定改参数数，也要传递给数据集迭代器，让数据集迭代器也知道此时需要提供refine模型的相关数据了
        if best_test < opt.refine_margin and not opt.refine_start:
            opt.refine_start = True
            opt.batch_size = int(opt.batch_size / opt.iteration)
            optimizer = optim.Adam(refiner.parameters(), lr=opt.lr)

            if opt.dataset == 'ycb':
                dataset = PoseDataset_ycb('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
            if opt.dataset == 'warehouse':
                dataset = PoseDataset_warehouse('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
            elif opt.dataset == 'linemod':
                dataset = PoseDataset_linemod('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
            dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
            if opt.dataset == 'ycb':
                test_dataset = PoseDataset_ycb('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
            if opt.dataset == 'warehouse':
                test_dataset = PoseDataset_warehouse('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
            elif opt.dataset == 'linemod':
                test_dataset = PoseDataset_linemod('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
            testdataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=opt.workers)
            
            opt.sym_list = dataset.get_sym_list()
            opt.num_points_mesh = dataset.get_num_points_mesh()

            print('>>>>>>>>----------Dataset loaded!---------<<<<<<<<\nlength of the training set: {0}\nlength of the testing set: {1}\nnumber of sample points on mesh: {2}\nsymmetry object list: {3}'.format(len(dataset), len(test_dataset), opt.num_points_mesh, opt.sym_list))

            criterion = Loss(opt.num_points_mesh, opt.sym_list)
            criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)

if __name__ == '__main__':
    main()

思路总结
大家也不要觉得太复杂，其上总的来说可以分为以下几个部分

1.搭建网络
	estimator = PoseNet(num_points = opt.num_points, num_obj = opt.num_objects)
	refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_objects)

2.构建训练验证测试数据集的迭代器：
	PoseDataset_ycb   PoseDataset_warehouse  PoseDataset_linemod


3.循环迭代，在迭代过程总，每个epoch结束之后，都会对当前模型进行一次判断
	if best_test < opt.refine_margin and not opt.refine_start:
	如果达到了要求，则开始训练refiner模型

其上的PoseNet网络，对应论文中如下部分：

是的，你没有看错，都包揽在其中，后续我们会大家一一分析讲解每个模块。其后的PoseRefineNet模块，对应论文的中的如下部分：

后续我也会为大家一一分析。

思路带领
那么，知道了网络框架，我们接下来要了解什么呢？当然是数据集和数据预处理，那么下小节我们就开始讲解，因为只有这样，我们才能知道网络的输入是什么。

2. 数据集讲解

通过前面，我们已经了解整个网络框架的训练流程，现在我们先来看看我们训练的数据集Linemod_preprocessed，本人存在如下上个子文件夹：

如果你的和我的长得不一样，那么抱歉，我也无能无力了。首先我们要了解的是，Linemod_preprocessed这个数据集中，一共用15种目标物体，首先我们来看看其中的data目录，如下：

分别对应十五个目标物体，随便既然一个文件夹，本人进入01如下：

其中以及划分了验证和训练数据，保存的都是对应图片的索引，这样固定的划分好，是为了方便论文的复现。下面是对每个文件的总结：

1.  depth（目录）:  深度图

2.  mask（目录）：目标物体的掩码，为标准的分割结果

3.  rgb（目录）：保存为RGB图像

5.  gt.yml（文件）：保存了拍摄每张图片时，其对应的旋转矩阵和偏移矩阵，以及目标物体的标准box，
    和该图像中目标物体所属于的类别

6.  info.yml（文件）：拍摄每张图像，对应的摄像头的内参，以及深度缩放的比例

这样，我们对于data已经了解完成了，我们再来看看Linemod_preprocessed/segnet_results，进入该文件如下：

其和Linemod_preprocessed/data/mask可以说时一一对应的，大小认真的对比下，发现他们对应的图片时相差无几的，那么这个又什么用？首先我们从名字来看，可以知道segnet_results时通过对目标物体进行语义分割得到的结果。再训练我们模型的时候，我们肯定希望时拿最标准的数据去训练我们的模型，所以训练和验证，都是使用最标准的mask（Linemod_preprocessed/data/mask），作为分割的标签。但是再测试的时候，是和应用场景达搭钩了，我们希望越接近实际场景越好，在实际场景中，DenseFusion网络是对分割出来的目标进行姿态预测，是的，是分割出来的物体，所以实际应用中，还需要一个分割网络。所以segnet_results中，保存的就是语义分割网络分割出来的图片，为了和实际场景更加接近，这也是倒是其和标准的/data/mask有点点区别的原因。

再下来我们就是介绍Linemod_preprocessed/models文件夹，进入如下：

这里的models可以理解为模型，是什么模型呢？那就是我们目标物体对应的模型，这些模型都用点云数据表示，并保存了起来。那么问题来了。在Linemod_preprocessed/data/xx/rgb中，有那么多的图像，这里的models是对应那个图像的模型呢？ 进入obj_01.ply文件可以看到如下：

ply
format ascii 1.0
comment VCGLIB generated
element vertex 5841
property float x
property float y
property float z
property float nx
property float ny
property float nz
property uchar red
property uchar green
property uchar blue
property uchar alpha
element face 11678
property list uchar int vertex_indices
end_header
-23.0565 17.4191 -6.756 -0.617876 0.77406 -0.138061 148 57 57 255 
-12.8773 -29.5554 -27.7156 -0.4007 -0.894836 0.196745 140 53 51 255 
-32.2776 -17.7052 -32.4505 -0.813037 -0.537012 -0.224921 151 61 58 255 

这些都是点云数据，简单的介绍可以参考如下连接：
https://blog.csdn.net/qq_36559293/article/details/90949295
可以想象得到，我们在对一个物体构建3D点云数据得时候，我们需要一个基准参考面，如果没有一个基准，我们拍摄多张照片，每个照片构建出来的点云数据都都不一样（那怕他们是同一个物体）。

所以Linemod_preprocessed\models中每个ply文件，其中的点云信息，都是以linemod\Linemod_preprocessed\data\xx\rgb\0000.png作为参考面的。有了这个参考面的点云数据（包含了整个目标的完整点云），我们就能根据其他的拍摄照片时的摄像头参数（旋转和偏移矩阵），计算出其他的视觉对应的点云数据了。

总的来说，我们对第一帧图片进行3D点云数据的建模（当然，整个过程需要多个面的图片-数据的制作过程），建模之后，我们对着目标随便拍摄一张图片，就能根据摄像头的参数，结合第一帧图片的3D点云数据，计算出该图片对应的3D点云数据了。

最后还有就是models_info.yml文件，其保存的是每个目标点云模型的半径，x，y，z轴的起始值和大小范围。

这样，我们对数据集的介绍就完成了。下面我们来看看对数据集的预处理过程

数据预处理
在tools/train.py中，我们可以看到如下：

PoseDataset_ycb()
PoseDataset_warehouse()
PoseDataset_linemod()

其都是对不同数据集的加载过程，对于本人要讲解的当然是PoseDataset_linemod，其代码的实现是在datasets/linemod/dataset.py：

import torch.utils.data as data
from PIL import Image
import os
import os.path
import errno
import torch
import json
import codecs
import numpy as np
import sys
import torchvision.transforms as transforms
import argparse
import json
import time
import random
import numpy
import numpy.ma as ma
import copy
import scipy.misc
import scipy.io as scio
import yaml
import cv2
numpy.set_printoptions(threshold= 1e6)

class PoseDataset(data.Dataset):
    def __init__(self, mode, num, add_noise, root, noise_trans, refine):
        """

        :param mode: 可以选择train，test，eval
        :param num: mesh点的数目
        :param add_noise:是否加入噪声
        :param root:数据集的根目录
        :param noise_trans:噪声增强的相关参数
        :param refine:是否需要为refine模型提供相应的数据
        """

        # 这里表示目标物体类别序列号
        self.objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
        # 可以选择train，test，eval
        self.mode = mode
        # 存储RGB图像的路径
        self.list_rgb = []
        # 存储深度图像的路径
        self.list_depth = []
        # 存储语音分割出来物体对应的mask
        self.list_label = []

        # 两个拼接起来,可以知道每张图片的路径，及物体类别，和图片下标
        self.list_obj = []
        self.list_rank = []

        # 矩阵信息，拍摄图片时的旋转矩阵和偏移矩阵，以及物体box
        self.meta = {}

        # 保存目标模型models点云数据，及models/obj_xx.ply文件中的数据
        self.pt = {}

        # 数据的所在的目录
        self.root = root
        # 噪声相关参数
        self.noise_trans = noise_trans
        self.refine = refine

        item_count = 0
        # 对每个目标物体的相关数据都进行处理
        for item in self.objlist:
            # 根据训练或者测试获得相应文件中的txt内容，其中保存的都是图片对应的名称数目
            if self.mode == 'train':
                input_file = open('{0}/data/{1}/train.txt'.format(self.root, '%02d' % item))
            else:
                input_file = open('{0}/data/{1}/test.txt'.format(self.root, '%02d' % item))

            # 循环txt记录的每张图片进行处理
            while 1:
                # 记录处理的数目
                item_count += 1
                input_line = input_file.readline()
                # test模式下，图片序列为10的倍数则continue
                if self.mode == 'test' and item_count % 10 != 0:
                    continue
                # 文件读取完成
                if not input_line:
                    break

                if input_line[-1:] == '\n':
                    input_line = input_line[:-1]
                # 把RGB图像的路径加载到self.list_rgb列表中
                self.list_rgb.append('{0}/data/{1}/rgb/{2}.png'.format(self.root, '%02d' % item, input_line))
                # 把data/x/depth图像的路径加载到self.list_depth列表中
                self.list_depth.append('{0}/data/{1}/depth/{2}.png'.format(self.root, '%02d' % item, input_line))

                # 如果是评估模式，则添加segnet_results图片的mask，否则添加data中的mask图片（该为标准mask）
                # 大家可以想想，训练的时候，肯定使用最标准的mask，但是在测试的时候，是要结合实际了，所以使用的
                # 是通过分割网络分割出来的mask，即则添加segnet_results中的图片
                if self.mode == 'eval':
                    self.list_label.append('{0}/segnet_results/{1}_label/{2}_label.png'.format(self.root, '%02d' % item, input_line))
                else:
                    self.list_label.append('{0}/data/{1}/mask/{2}.png'.format(self.root, '%02d' % item, input_line))

                # 把物体的下标，和txt读取到的图片数目标记分别添加到list_obj,list_rank中
                self.list_obj.append(item)
                self.list_rank.append(int(input_line))

            # gt.yml主要保存的，是拍摄图片时，物体的旋转矩阵以及偏移矩阵，以及物体标签的box
            # 有了该参数，我们就能把对应的图片，从2维空间恢复到3维空间了
            meta_file = open('{0}/data/{1}/gt.yml'.format(self.root, '%02d' % item), 'r')
            self.meta[item] = yaml.load(meta_file)

            # 这里保存的是目标物体，拍摄物体第一帧的点云数据，可以成为模型数据
            self.pt[item] = ply_vtx('{0}/models/obj_{1}.ply'.format(self.root, '%02d' % item))
            
            print("Object {0} buffer loaded".format(item))

        # 读取所有图片的路径，以及文件信息之后，打印图片数目
        self.length = len(self.list_rgb)

        # 摄像头的中中心坐标
        self.cam_cx = 325.26110
        self.cam_cy = 242.04899
        # 摄像头的x，y轴的长度
        self.cam_fx = 572.41140
        self.cam_fy = 573.57043

        # 列举处图片的x和y坐标
        # xmap[480,640]   ymap[480,640]
        self.xmap = np.array([[j for i in range(640)] for j in range(480)])
        #print(self.xmap.shape)
        self.ymap = np.array([[i for i in range(640)] for j in range(480)])

        # 设定获取目标物体点云的数据
        self.num = num
        self.add_noise = add_noise
        self.trancolor = transforms.ColorJitter(0.2, 0.2, 0.2, 0.05)
        self.norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

        # 边界列表，可以想象把一个图片切割成了多个坐标
        self.border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]

        # 点云的最大和最小数目
        self.num_pt_mesh_large = 500
        self.num_pt_mesh_small = 500

        # 标记对称物体的序列号
        self.symmetry_obj_idx = [7, 8]

    # 迭代的时候，会调用该函数
    def __getitem__(self, index):
        # 根据索引获得对应图像的RGB像素
        img = Image.open(self.list_rgb[index])
        #ori_img = np.array(img)
        # 根据索引获得对应图像的深度图像素
        depth = np.array(Image.open(self.list_depth[index]))
        # 根据索引获得对应图像的mask像素
        label = np.array(Image.open(self.list_label[index]))
        # 获得物体属于的类别的序列号
        obj = self.list_obj[index]
        # 获得该张图片物体图像的下标
        rank = self.list_rank[index]        

        # 如果该目标物体的序列为2，暂时不对序列为2的物体图像进行处理
        if obj == 2:
            # 对该物体的每个图片下标进行循环
            for i in range(0, len(self.meta[obj][rank])):
                # 验证该图片目标十分为2，如果是则赋值给meta
                if self.meta[obj][rank][i]['obj_id'] == 2:
                    meta = self.meta[obj][rank][i]
                    break
        # 如果目标物体的序列不为2，获得目标物体第一帧图片的旋转以及偏移矩阵
        else:
            meta = self.meta[obj][rank][0]

        # 只要像素不为0的，返回值为Ture，如果为0的像素，还是返回False
        mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0))
        #print(mask_depth)

        # 标准的数据中的mask是3通道的，通过网络分割出来的mask，其是单通道的
        if self.mode == 'eval':
            mask_label = ma.getmaskarray(ma.masked_equal(label, np.array(255)))
        else:
            mask_label = ma.getmaskarray(ma.masked_equal(label, np.array([255, 255, 255])))[:, :, 0]

        # 把mask和深度图结合到到一起，物体存在的区域像素为True，背景像素为False
        mask = mask_label * mask_depth
        # print('1='*50)
        # print(mask.shape)
        # print(mask)
        # p_img = np.transpose(mask, (1, 2, 0))
        # scipy.misc.imsave('evaluation_result/{0}_mask.png'.format(index), p_img)

        # 对图像加入噪声
        if self.add_noise:
            img = self.trancolor(img)

        # [b,h,w,c] --> [b,c,h,w]
        img = np.array(img)[:, :, :3]
        img = np.transpose(img, (2, 0, 1))
        img_masked = img

        # 如果为eval模式，根据mask_label获得目标的box（rmin-rmax表示行的位置，cmin-cmax表示列的位置）
        if self.mode == 'eval':
            # mask_to_bbox表示的根据mask生成合适的box
            rmin, rmax, cmin, cmax = get_bbox(mask_to_bbox(mask_label))
        # 如果不是eval模式，则从gt.yml文件中，获取最标准的box
        else:
            rmin, rmax, cmin, cmax = get_bbox(meta['obj_bb'])

        # 根据确定的行和列，对图像进行截取，就是截取处包含了目标物体的图像
        img_masked = img_masked[:, rmin:rmax, cmin:cmax]
        #p_img = np.transpose(img_masked, (1, 2, 0))
        #scipy.misc.imsave('evaluation_result/{0}_input.png'.format(index), p_img)

        # 获得目标物体旋转矩阵的参数，以及偏移矩阵参数
        target_r = np.resize(np.array(meta['cam_R_m2c']), (3, 3))
        target_t = np.array(meta['cam_t_m2c'])

        # 对偏移矩阵添加噪声
        add_t = np.array([random.uniform(-self.noise_trans, self.noise_trans) for i in range(3)])

        # 对mask图片的目标部分进行剪裁，变成拉平，变成一维
        choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0]
        # print('2='*50)
        # print(choose.shape)
        # print(choose)


        # 如果剪切下来的部分面积为0，则直接返回五个0的元组，即表示没有目标物体
        if len(choose) == 0:
            cc = torch.LongTensor([0])
            return(cc, cc, cc, cc, cc, cc)

        # 如果剪切下来目标图片的像素，大于了点云的数目（一般都是这种情况）
        if len(choose) > self.num:
            # c_mask全部设置为0，大小和choose相同
            c_mask = np.zeros(len(choose), dtype=int)

            # 前self.num设置为1
            c_mask[:self.num] = 1

            # 随机打乱
            np.random.shuffle(c_mask)

            # 选择c_mask不是0的部分，也就是说，只选择了500个像素，注意nonzero()返回的是索引
            choose = choose[c_mask.nonzero()]
        # 如果剪切像素点的数目，小于了点云的数目
        else:
            # 这使用0填补，调整到和点云数目一样大小（500）
            choose = np.pad(choose, (0, self.num - len(choose)), 'wrap')

        # 得到的choose输出为（500,）这里保存的,是挑选出来的500个索引
        #print(choose.shape)
        #print(choose)

        # 把深度图，对应着物体的部分也剪切下来，然后拉平，变成一维，挑选坐标
        depth_masked = depth[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)

        # 把物体存在于原图的位置坐标剪切下来，拉平，然后进行挑选坐标挑选
        xmap_masked = self.xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)
        ymap_masked = self.ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32)

        choose = np.array([choose])

        # 摄像头缩放参数
        cam_scale = 1.0
        pt2 = depth_masked / cam_scale

        # 为对坐标进行正则化做准备
        pt0 = (ymap_masked - self.cam_cx) * pt2 / self.cam_fx
        pt1 = (xmap_masked - self.cam_cy) * pt2 / self.cam_fy

        # 把y，x，depth他们3个坐标合并在一起，变成点云数据，
        cloud = np.concatenate((pt0, pt1, pt2), axis=1)
        #print(cloud)

        # 这里把点云数据除以1000，是为了根据深度进行正则化
        cloud = cloud / 1000.0

        #print(cloud)
        # 对点云添加噪声
        if self.add_noise:
            cloud = np.add(cloud, add_t)

        #fw = open('evaluation_result/{0}_cld.xyz'.format(index), 'w')
        #for it in cloud:
        #    fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
        #fw.close()


        # 存储在obj_xx.ply中的点云数据，对其进行正则化，也就是目标物体的点云信息
        model_points = self.pt[obj] / 1000.0

        # 随机删除多余的点云数据，训练时，只需要num_pt_mesh_small数目的点云
        dellist = [j for j in range(0, len(model_points))]
        dellist = random.sample(dellist, len(model_points) - self.num_pt_mesh_small)
        model_points = np.delete(model_points, dellist, axis=0)

        #fw = open('evaluation_result/{0}_model_points.xyz'.format(index), 'w')
        #for it in model_points:
        #    fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
        #fw.close()

        # 根据model_points（第一帧目标模型对应的点云信息），以及target（目前迭代这张图片）的旋转和偏移矩阵，计算出对应的点云数据
        target = np.dot(model_points, target_r.T)

        if self.add_noise:
            target = np.add(target, target_t / 1000.0 + add_t)
            out_t = target_t / 1000.0 + add_t
        else:
            target = np.add(target, target_t / 1000.0)
            out_t = target_t / 1000.0

        #fw = open('evaluation_result/{0}_tar.xyz'.format(index), 'w')
        #for it in target:
        #    fw.write('{0} {1} {2}\n'.format(it[0], it[1], it[2]))
        #fw.close()


        # 总结：
        # cloud：由深度图计算出来的点云，该点云数据以本摄像头为参考坐标
        # choose：所选择点云的索引
        # img_masked：通过box剪切下来的RGB图像
        # target：根据model_points点云信息，以及旋转偏移矩阵转换过的点云信息
        # model_points：目标初始帧（模型）对应的点云信息
        # [self.objlist.index(obj)]：目标物体的序列编号
        return torch.from_numpy(cloud.astype(np.float32)), \
               torch.LongTensor(choose.astype(np.int32)), \
               self.norm(torch.from_numpy(img_masked.astype(np.float32))), \
               torch.from_numpy(target.astype(np.float32)), \
               torch.from_numpy(model_points.astype(np.float32)), \
               torch.LongTensor([self.objlist.index(obj)])

    def __len__(self):
        return self.length

    def get_sym_list(self):
        return self.symmetry_obj_idx

    def get_num_points_mesh(self):
        if self.refine:
            return self.num_pt_mesh_large
        else:
            return self.num_pt_mesh_small



border_list = [-1, 40, 80, 120, 160, 200, 240, 280, 320, 360, 400, 440, 480, 520, 560, 600, 640, 680]
img_width = 480
img_length = 640


def mask_to_bbox(mask):
    mask = mask.astype(np.uint8)
    #_, contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    x = 0
    y = 0
    w = 0
    h = 0
    for contour in contours:
        tmp_x, tmp_y, tmp_w, tmp_h = cv2.boundingRect(contour)
        if tmp_w * tmp_h > w * h:
            x = tmp_x
            y = tmp_y
            w = tmp_w
            h = tmp_h
    return [x, y, w, h]


def get_bbox(bbox):
    bbx = [bbox[1], bbox[1] + bbox[3], bbox[0], bbox[0] + bbox[2]]
    if bbx[0] < 0:
        bbx[0] = 0
    if bbx[1] >= 480:
        bbx[1] = 479
    if bbx[2] < 0:
        bbx[2] = 0
    if bbx[3] >= 640:
        bbx[3] = 639                
    rmin, rmax, cmin, cmax = bbx[0], bbx[1], bbx[2], bbx[3]
    r_b = rmax - rmin
    for tt in range(len(border_list)):
        if r_b > border_list[tt] and r_b < border_list[tt + 1]:
            r_b = border_list[tt + 1]
            break
    c_b = cmax - cmin
    for tt in range(len(border_list)):
        if c_b > border_list[tt] and c_b < border_list[tt + 1]:
            c_b = border_list[tt + 1]
            break
    center = [int((rmin + rmax) / 2), int((cmin + cmax) / 2)]
    rmin = center[0] - int(r_b / 2)
    rmax = center[0] + int(r_b / 2)
    cmin = center[1] - int(c_b / 2)
    cmax = center[1] + int(c_b / 2)
    if rmin < 0:
        delt = -rmin
        rmin = 0
        rmax += delt
    if cmin < 0:
        delt = -cmin
        cmin = 0
        cmax += delt
    if rmax > 480:
        delt = rmax - 480
        rmax = 480
        rmin -= delt
    if cmax > 640:
        delt = cmax - 640
        cmax = 640
        cmin -= delt
    return rmin, rmax, cmin, cmax


def ply_vtx(path):
    f = open(path)
    assert f.readline().strip() == "ply"
    f.readline()
    f.readline()
    N = int(f.readline().split()[-1])
    while f.readline().strip() != "end_header":
        continue
    pts = []
    for _ in range(N):
        pts.append(np.float32(f.readline().split()[:3]))
    return np.array(pts)

我相信，注释已经很详细，我们回到tools/train.py:

	points, choose, img, target, model_points, idx = data
	# points：由深度图计算出来的点云，该点云数据以摄像头为参考坐标
	# choose：所选择点云的索引，
	# img：通过box剪切下来的RGB图像
	# target：根据model_points点云信息，以及旋转偏移矩阵转换过的点云信息
	# model_points：目标初始帧（模型）对应的点云信息
	# idx：目标物体的序列编号

其上可以说，就是对整个数据加载，以及预测处理的目的了。

细心的朋友应该发现了，修改batch_size参数，似乎没有什么用，其主要的原因是torch.utils.data.DataLoader()函数的参数batch_size=1设置为1，集每次都是获得一个样本，如果强行修改为其他的整数（如本人修改为2），会报错如下：

taloader.py", line 209, in default_collate
    return torch.stack(batch, 0, out=out)
RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 160 and 120 in dimension 2 at /pytorch/aten/src/TH/generic/THTensorMoreMath.cpp:1307

大概的原因就是，对于每个我们分割出来的目标图像，其大小都是不一样的，所以我们只能一张一张的去进行训练。多个图片大小不一样的图像，没有办法统一成一个batch。

3. PoseNet姿态估算网络详解

网络输入

通过上小结，我们可以知道在tools/train.py可以看到本人的如下注释：

for i, data in enumerate(dataloader, 0):
    points, choose, img, target, model_points, idx = data
    # points：由深度图计算出来的点云，该点云数据以摄像头主轴参考坐标
    # choose：所选择点云的索引，[bs, 1, 500]
    # img：通过box剪切下来的RGB图像
    # target：根据model_points点云信息，以及旋转偏移矩阵转换过的点云信息[bs,500,3]
    # model_points：目标初始帧（模型）对应的点云信息[bs,500,3]
    # idx：目标物体的序列编号
	pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)

这个就是我们整个网络的输入数据了，以及对应的PoseNet（其上的estimator就是），那么这些数据在整个网络中是怎么运转的呢？首先我们来看看PoseNet网络的实现，其代码在lib/network.py之中，这是本人对class PoseNet(nn.Module)的注解（大致浏览下，后面有更加详细的带读）

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
       
class PoseNet(nn.Module):
    def __init__(self, num_points, num_obj):
        """
        :param num_points:输入网络的点云数目
        :param num_obj: 目标物体的种类
        """
        # 调用父类的初始化函数，不知道有什么用
        super(PoseNet, self).__init__()

        # 点云的数目
        self.num_points = num_points

        # 做过修改的Resnet
        self.cnn = ModifiedResnet()
        summary(self.cnn,(3, 120, 120))

        self.feat = PoseNetFeat(num_points)
        summary(self.feat.cuda(), [(3, 500),(32,500)])

        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):
        """
        PoseNet的前向传播，进行姿态预测
        :param img: RGB图像的像素[bs,3,h,w]
        :param x: 点云数据[bs, 500, 3]
        :param choose: 选择点云的index下标[bs, 1, 500]
        :param obj: 目标物体的序列号[bs, 1]
        :return:对图像预测的姿态
        """

        #print('img.shape: {0}',format(img.shape))
        #print('x.shape: {0}', format(x.shape))
        #print('choose.shape: {0}', format(choose.shape))
        #print('obj.shape: {0}', format(obj.shape))

        #out_img[bs, 32, h, w]
        out_img = self.cnn(img)

        bs, di, _, _ = out_img.size()

        # 进行resize操作变成
        emb = out_img.view(bs, di, -1)

        # 进行复制，复制di=32次，[1, 32, 500]
        choose = choose.repeat(1, di, 1)
        #print(choose.shape)

        # contiguous()是为了保证在GPU上连续分配
        # 并且针对emb每个通道（di），选取500个像素，及[bs, di=32, ?]-->[1, di=32, 500]
        emb = torch.gather(emb, 2, choose).contiguous()
        #print(emb.shape)

        # [bs, 500, 3]-->[bs,3,500]
        x = x.transpose(2, 1).contiguous()

        # x: [bs, 3, 500]，后续通过卷积变成论文中的geometry embeddings
        # emb: [bs, 32, 500],论文中的color embeddings
        #  ap_x[bs, 1408 = 128 + 256 + 1024, 500]
        ap_x = self.feat(x, emb)
        #print('ap_x.shape: {0}', format(ap_x.shape))

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

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

        rx = F.relu(self.conv3_r(rx))
        tx = F.relu(self.conv3_t(tx))
        cx = F.relu(self.conv3_c(cx))


        # rx.shape: [bs, num_obj, 4, 500]
        rx = self.conv4_r(rx).view(bs, self.num_obj, 4, self.num_points)
        # tx.shape: [bs, num_obj, 3, 500]
        tx = self.conv4_t(tx).view(bs, self.num_obj, 3, self.num_points)
        # cx.shape: [bs, num_obj, 1, 500]
        cx = torch.sigmoid(self.conv4_c(cx)).view(bs, self.num_obj, 1, self.num_points)
        # print('='*50)
        # print('rx.shape: {0}', format(rx.shape))
        # print('tx.shape: {0}', format(tx.shape))
        # print('cx.shape: {0}', format(cx.shape))

        # 选择预测对应目标内标的矩阵参数
        b = 0

        #[bs, 4, 500]
        out_rx = torch.index_select(rx[b], 0, obj[b])
        #[bs, 3, 500]
        out_tx = torch.index_select(tx[b], 0, obj[b])
        #[bs, 1, 500]
        out_cx = torch.index_select(cx[b], 0, obj[b])
        # print('out_rx.shape: {0}', format(out_rx.shape))
        # print('out_tx.shape: {0}', format(out_tx.shape))
        # print('out_cx.shape: {0}', format(out_cx.shape))

        # [bs, 500, 4]
        out_rx = out_rx.contiguous().transpose(2, 1).contiguous()
        # [bs, 500, 3]
        out_cx = out_cx.contiguous().transpose(2, 1).contiguous()
        # [bs, 500, 1]
        out_tx = out_tx.contiguous().transpose(2, 1).contiguous()

        # detach()表示从图中分离出来，不做反向传播
        return out_rx, out_tx, out_cx, emb.detach()

源码带读

首先我们从class PoseNet(nn.Module)这个类看起。直接进入其def forward(self, img, x, choose, obj)函数，这是网络的前线传播。可以看到如下代码：

	#out_img[bs, 32, h, w]
	out_img = self.cnn(img)

其中的CNN对应该类初始化函数中的：

    # 做过修改的Resnet
    self.cnn = ModifiedResnet()
    summary(self.cnn,(3, 120, 120))

想看网络结构的朋友，把 summary(self.cnn,(3, 120, 120))添加上去，就能看到整个网络的机构了，我这里就不复制了。总的来说，对应论文Figure 2如下部分：

其中的image crop就是forward(self, img, x, choose, obj)参数中的img，self.cnn(img)网络的输出out_img或者后面的emb就图示中的color embeddings。这里呢，大家只要知道网络输入img[bs,3,H,W]然后输出emb[bs,32,H,W]，也就是说，每个像素都被抽象成了一个32维度的向量。

然后大家要注意的是forward函数输入x的大小为[bs, 500, 3]，也就是说，在数据预处理的时候，我们已经选择了500个点云进来。但是我们图片的特征向量[bs,32,H,W]可是有HxW个像素，论文中说到color embeddings和geometry embeddings要进行稠密（像素级别）的融合，那你也得像素和点云数目一样多，才能融合吧，不然融合个锤锤啊。

所以呢，这里就借助传入的choose参数，选择了和点云对应位置的color embeddings，这样他们的大小都为500了，就能进行融合了。融合呢，融合过程调用如下函数：

    # x: [bs, 3, 500]，后续通过卷积变成论文中的geometry embeddings
    # emb: [bs, 32, 500],论文中的color embeddings
    #  ap_x[bs, 1408 = 128 + 256 + 1024, 500]
    ap_x = self.feat(x, emb)
    #print('ap_x.shape: {0}', format(ap_x.shape))

可以知道，这个网络的输入点云x[bs, 3, 500]（空间信息）和emb: [bs, 32, 500]（颜色信息）。合成之后呢，每个像素与点云都融合成了1408 = 128 + 256 + 1024维度的向量，那么这些都是什么东西？先看self.feat结构的实现

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):
        """
        :param x: [bs, 3, 500]，点云数据
        :param emb: [bs, 32, 500],论文中的color embeddings
        """
        # [bs, 3, 500]-->[bs,64,500]，可以理解为论文中的论文中的geometry embeddings
        x = F.relu(self.conv1(x))
        # [bs, 32, 500]-->[bs,64,500]
        emb = F.relu(self.e_conv1(emb))
        # [bs,128,500]
        pointfeat_1 = torch.cat((x, emb), dim=1)
        #print('pointfeat_1.shape: {0}', format(pointfeat_1.shape))

        #[bs, 64, 500]-->[bs, 128, 500]
        x = F.relu(self.conv2(x))
        # [bs, 64, 500]-->[bs, 128, 500]
        emb = F.relu(self.e_conv2(emb))
        # [bs, 256, 500]
        pointfeat_2 = torch.cat((x, emb), dim=1)
        #print('pointfeat_2.shape: {0}', format(pointfeat_2.shape))

        # [bs, 256, 500]-->[bs, 512, 500]
        x = F.relu(self.conv5(pointfeat_2))
        # [bs, 512, 500]-->[bs, 1024, 500]
        x = F.relu(self.conv6(x))
        #print('x.shape: {0}', format(x.shape))

        ap_x = self.ap1(x)

        # ap_x[bs, 1024, 500]
        ap_x = ap_x.view(-1, 1024, 1).repeat(1, 1, self.num_points)
        #print('ap_x.shape: {0}', format(ap_x.shape))

        return torch.cat([pointfeat_1, pointfeat_2, ap_x], 1) #128 + 256 + 1024

可以明显的看到，网络的返回注释：

return torch.cat([pointfeat_1, pointfeat_2, ap_x], 1) #128 + 256 + 1024

其500个中的每个128 + 256 + 1024对应分别如下

大家可能奇怪，为什么要使用这样的3段拼接起来，首先我们来看看pointfeat_1：

        # [bs, 3, 500]-->[bs,64,500]
        x = F.relu(self.conv1(x))
        # [bs, 32, 500]-->[bs,64,500]
        emb = F.relu(self.e_conv1(emb))
        # [bs,128,500]
        pointfeat_1 = torch.cat((x, emb), dim=1)
        #print('pointfeat_1.shape: {0}', format(pointfeat_1.shape))

其为论文如下部分：

红色框表示他们融合的过程，绿色的框表示如何之后，直接当做pixel-wise feature的一部分。为什么要这么做呢？首先来说，他们是像素级别的融合，融合之后没有做任何出来，也就是说，现在他们保留了每个像素的空间和几何信息。但是这里出现了问题，我们在姿态估算的时候，出来要像素级别的信息，我们应该也还需要全局的信息吧，去结合上下文的联系，所以呢，后面就出现了pointfeat_2，ap_x。他们都是在前面的基础上，进行了特征提取，pointfeat_2属于中间部分，包含的信息应该是点与点之间的联系。ap_x存储的就是全局的信息了。

这样分析下来，是不是所谓的ap_x = self.feat(x, emb)得到的ap_x就很有意义了，其并不是乱七八糟的融合。而是包含了全局，局部，和像素级别的信息。拿到这么完美的一个东西，当然就可以直接做姿态预测了，所以后面再代码中，我们就能看到：

        #[bs, 4, 500]
        out_rx = torch.index_select(rx[b], 0, obj[b])
        #[bs, 3, 500]
        out_tx = torch.index_select(tx[b], 0, obj[b])
        #[bs, 1, 500]，置信度
        out_cx = torch.index_select(cx[b], 0, obj[b])

但是这里大家要注意一下，这里预测的是姿态的旋转参数和偏移参数，其中的out_rx并不是偏移矩阵。还有就是out_cx表示的是置信度，以及其中的500，表示的含义是对每个像素都有预测，后面我们就要从这500个预测中，选出最好的一个结果最终的预测结果。

到这里PoseNet的前向传播基本就完成了，下小结我们来看看他的损失函数。

4. PoseNet网络loss详解（重点篇）

代码引导
在tools/train.py中，我们可以看到如下代码：

	 points, choose, img, target, model_points, idx = data
     # points：由深度图计算出来的点云，该点云数据以摄像头主轴参考坐标
     # choose：所选择点云的索引，[bs, 1, 500]
     # img：通过box剪切下来的RGB图像
     # target：根据model_points点云信息，以及旋转偏移矩阵转换过的点云信息[bs,500,3]
     # model_points：目标初始帧（模型）对应的点云信息[bs,500,3]
     # idx：目标物体的序列编号

	 pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)
	 # 进行预测获得，获得预测的姿态，姿态预测之前的特征向量
	 # pred_r: 预测的旋转参数[bs, 500, 4]
	 # pred_t: 预测的偏移参数[bs, 500, 3]
	 # pred_c: 预测的置信度[bs, 500, 1]
	
	
	 # 对结果进行评估，计算loss
	 loss, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start)

通过上小结，我们可以知道estimator是对姿态进行估算，其函数返回的emb就是当前需要估算姿态图像（RBG）抽取出来的特征向量。从data中迭代出来的target，已经从数据预处理章节中推导出来，他是由model_points（目标物体第一帧点云）根据标准的参数转换到当前帧的点云数据，可以理解为他就是一个标签ground truth。points表示的是当前帧的点云数据，注意target的参考基准为model_points，points的参考基准是摄像头。refine_start标记已经是否开始了refine网络的训练。好了这样大家就明白网络的输入了，我们进入lib/loss.py来看看其上criterion函数的实现（大致浏览一下，后面本人有代码领读的-麻烦的事情交给我就好，把你的时间用在刀刃上）。

from torch.nn.modules.loss import _Loss
from torch.autograd import Variable
import torch
import time
import numpy as np
import torch.nn as nn
import random
import torch.backends.cudnn as cudnn
from lib.knn.__init__ import KNearestNeighbor


def loss_calculation(pred_r, pred_t, pred_c, target, model_points, idx, points, w, refine, num_point_mesh, sym_list):
    """
    :param pred_r: 预测的旋转参数[bs, 500, 4]，相对于摄像头
    :param pred_t: 预测的偏移参数[bs, 500, 3]，相对于摄像头
    :param pred_c: 预测的置信度参数[bs, 500, 1]，相对于摄像头
    :param target: 目标姿态，也就是预测图片，通过标准偏移矩阵，结合model_points求得图片对应得点云数据[bs,500,3]，这里点云数据，就是学习的目标数据
    :param model_points:目标模型的点云数据-第一帧[bs,500,3]
    :param idx:随机训练的一个索引
    :param points:由深度图计算出来的点云，也就是说该点云数据以摄像头为参考坐标
    :param refine:标记是否已经开始训练refine网络
    :param num_point_mesh:500
    :param sym_list:对称模型的序列号
    """
    print('='*50)
    #print('pred_r.shape: {0}', format(pred_r.shape))
    # print('target.shape: {0}', format(target.shape))
    # print('model_points.shape: {0}', format(model_points.shape))
    # print('points.shape: {0}', format(points.shape))

    knn = KNearestNeighbor(1)
    # [bs, 500, 1]
    bs, num_p, _ = pred_c.size()

    # 把预测的旋转矩阵进行正则化
    pred_r = pred_r / (torch.norm(pred_r, dim=2).view(bs, num_p, 1))
    #print('pred_r.shape: {0}', format(pred_r.shape))

    # base[bs，500, 4] -->[500, 3, 3]，把预测的旋转参数，转化为旋转矩阵矩阵，
    base = torch.cat(((1.0 - 2.0*(pred_r[:, :, 2]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1),\
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] - 2.0*pred_r[:, :, 0]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 1]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 3]*pred_r[:, :, 0]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 3]**2)).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (-2.0*pred_r[:, :, 0]*pred_r[:, :, 2] + 2.0*pred_r[:, :, 1]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (2.0*pred_r[:, :, 0]*pred_r[:, :, 1] + 2.0*pred_r[:, :, 2]*pred_r[:, :, 3]).view(bs, num_p, 1), \
                      (1.0 - 2.0*(pred_r[:, :, 1]**2 + pred_r[:, :, 2]**2)).view(bs, num_p, 1)), dim=2).contiguous().view(bs * num_p, 3, 3)
    #print('base.shape: {0}', format(base.shape))

    # 把相对于摄像头的偏移矩阵记录下来
    ori_base = base

    # [3, 3, 500]
    base = base.contiguous().transpose(2, 1).contiguous()

    # 复制num_p=500次，[bs,1,500,3]-->[500,500,3]，这里的复制操作，主要是因为每个ground truth（target）点云，
    # 需要与所有的predicted点云做距离差，
    model_points = model_points.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    #print('model_points.shape: {0}', format(model_points.shape))

    # 复制num_p=500次，[bs,1,500,3]-->[500,500,3]，这里的复制操作，主要是因为每个ground truth（target）点云，
    # 需要与所有的predicted点云做距离差，
    target = target.view(bs, 1, num_point_mesh, 3).repeat(1, num_p, 1, 1).view(bs * num_p, num_point_mesh, 3)
    #print('target.shape: {0}', format(target.shape))


    # 把初始的目标点云（已经通过标准的pose进行了变换）记录下来
    ori_target = target
    pred_t = pred_t.contiguous().view(bs * num_p, 1, 3)
    # 把原始预测的偏移矩阵记录下来，这里的t是相对摄像头的
    ori_t = pred_t

    # 当前帧的点云，结合深度图计算而来，也就是说该点云信息是以摄像头为参考目标
    points = points.contiguous().view(bs * num_p, 1, 3)
    pred_c = pred_c.contiguous().view(bs * num_p)

    # 为批量矩阵相乘,model_points与旋转矩阵相乘加上偏移矩阵，得到当前帧对应的点云姿态，该点云的姿态是以model_points为参考的
    #pred[500,500,3]
    pred = torch.add(torch.bmm(model_points, base), points + pred_t)
    #print('pred.shape: {0}', format(pred.shape))

    #print('refine: {0}', format(refine))
    if not refine:
        # 如果是对称的物体
        if idx[0].item() in sym_list:

            # [500,500,3]-->[3,250000]
            target = target[0].transpose(1, 0).contiguous().view(3, -1)
            #print('target.shape: {0}', format(target.shape))

            # [500,500,3]-->[3,250000]
            pred = pred.permute(2, 0, 1).contiguous().view(3, -1)
            #print('pred.shape: {0}', format(pred.shape))

            # [1, 1, 250000]，target的每个点云和pred的所有点云进行对比，找到每个target点云与pred的所有点云，距离最近点云的索引（pred）
            inds = knn(target.unsqueeze(0), pred.unsqueeze(0))
            #print('inds.shape: {0}', format(inds.shape))

            # [3, 250000]，从target点云中，根据计算出来的min索引，全部挑选出来
            target = torch.index_select(target, 1, inds.view(-1).detach() - 1)
            #print('target.shape: {0}', format(target.shape))

            # [500, 500, 3]
            target = target.view(3, bs * num_p, num_point_mesh).permute(1, 2, 0).contiguous()

            # [500, 500, 3]
            pred = pred.view(3, bs * num_p, num_point_mesh).permute(1, 2, 0).contiguous()

    # 求得预测点云和目标点云的平均距离（每个点云）,按照论文，把置信度和点云距离关联起来
    dis = torch.mean(torch.norm((pred - target), dim=2), dim=1)
    loss = torch.mean((dis * pred_c - w * torch.log(pred_c)), dim=0)
    

    # 下面的操作都是为refine模型训练的准备
    pred_c = pred_c.view(bs, num_p)

    # which_max表示的索引下标，即找到置信度最高的哪个下标
    how_max, which_max = torch.max(pred_c, 1)
    print(how_max,which_max)
    #print('which_max.shape: {0}', format(which_max.shape))

    dis = dis.view(bs, num_p)

    # 获得最好的偏移矩阵，这里的t是相对model_points的
    t = ori_t[which_max[0]] + points[which_max[0]]
    points = points.view(1, bs * num_p, 3)

    # 求得500中置信度最高的旋转矩阵，相对于摄像头的
    ori_base = ori_base[which_max[0]].view(1, 3, 3).contiguous()
    ori_t = t.repeat(bs * num_p, 1).contiguous().view(1, bs * num_p, 3)

    # 根据预测最好的旋转矩阵，求得新的当前帧对应的点云，注意这里是一个减号的操作，并且其中的ori_t是相对于摄像头的
    # （但是ori_t和ori_base都是预测出来的，就是返回去肯定存在偏差的）
    new_points = torch.bmm((points - ori_t), ori_base).contiguous()

    new_target = ori_target[0].view(1, num_point_mesh, 3).contiguous()
    ori_t = t.repeat(num_point_mesh, 1).contiguous().view(1, num_point_mesh, 3)

    # 根据预测最好的旋转矩阵，求得新的当前帧对应的点云，注意这里是一个减号的操作，并且其中的ori_t是相对于摄像头的
    # （但是ori_t和ori_base都是预测出来的，就是返回去肯定存在偏差的，这里的偏差因为new_target是标准的，所以应该少一些）
    new_target = torch.bmm((new_target - ori_t), ori_base).contiguous()

    # print('------------> ', dis[0][which_max[0]].item(), pred_c[0][which_max[0]].item(), idx[0].item())
    del knn

    # loss：根据每个点云计算出来的平均loss
    # 对应预测置信度度最高，target与预测点云之间的最小距离
    # new_points：根据最好的旋转矩阵，求得当前帧的点云
    # new_target：就是根据model_points求得的标椎点云
    return loss, dis[0][which_max[0]], new_points.detach(), new_target.detach()


class Loss(_Loss):

    def __init__(self, num_points_mesh, sym_list):
        super(Loss, self).__init__(True)
        self.num_pt_mesh = num_points_mesh
        self.sym_list = sym_list

    def forward(self, pred_r, pred_t, pred_c, target, model_points, idx, points, w, refine):

        return loss_calculation(pred_r, pred_t, pred_c, target, model_points, idx, points, w, refine, self.num_pt_mesh, self.sym_list)

代码领读
1.首先我们我们来说说knn = KNearestNeighbor(1)，他作用其实很简单，就是为了加快计算的速度，所以作者使用的是编译成的C++库，看了我论文介绍的朋友，大家应该就知道，在训练对称物体的时候，每个target点云，需要在由估算姿态pose转化而来的pred cloud所有点云中，找到和他距离最近的哪个点云，然后把所有的距离加起来求得平均距离，当做loss。

2.在论文解读中，我多次强调，其是一个稠密(像素/点云级别)的姿态估算，也就是说，他会对每个点云/像素都做一个估算，并且都带有其相应的置信度。所以能够看到这样如下代码：

    # which_max表示的索引下标，即找到置信度最高的哪个下标
    how_max, which_max = torch.max(pred_c, 1)
    
    # 获得最好的偏移矩阵
    t = ori_t[which_max[0]] + points[which_max[0]]
    points = points.view(1, bs * num_p, 3)    

    # 求得500中置信度最高的旋转矩阵
    ori_base = ori_base[which_max[0]].view(1, 3, 3).contiguous()
    ori_t = t.repeat(bs * num_p, 1).contiguous().view(1, bs * num_p, 3)

这样我们就能找到最好的旋转和偏移矩阵了。

3.重难点3.重难点，这里就是思维的转折点了，如果没有认真看我论文翻译的（你在原论文中，是看不到这样详细的讲解的），请通过下面链接翻到后面的部分，读一下：
姿态估计0-03：DenseFusion(6D姿态估计)-白话给你讲论文-翻译无死角（1）
读了之后，大家再来看下面的代码：

    # 根据预测最好的旋转矩阵，求得新的当前帧对应的点云，注意这里是一个减号的操作，并且其中的ori_t是相对于摄像头的
    # （但是ori_t和ori_base都是预测出来的，就是返回去肯定存在偏差的）
    new_points = torch.bmm((points - ori_t), ori_base).contiguous()


    # 根据预测最好的旋转矩阵，求得新的当前帧对应的点云，注意这里是一个减号的操作，并且其中的ori_t是相对于摄像头的
    # （但是ori_t和ori_base都是预测出来的，就是返回去肯定存在偏差的，这里的偏差因为new_target是标准的，所以应该少一些）
    new_target = torch.bmm((new_target - ori_t), ori_base).contiguous()

5. PoseRefineNet网络与loss详解（重点篇）

代码引导

没有办法，还是得把代码贴上来一下tools/train.py：

    # 进行预测获得，获得预测的姿态，姿态预测之前的特征向量
    # pred_r: 预测的旋转参数[bs, 500, 4]
    # pred_t: 预测的偏移参数[bs, 500, 3]
    # pred_c: 预测的置信度[bs, 500, 1]，置信度
    #
    pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)

    # 对结果进行评估，计算loss
    loss, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start)

    # 如果已经对refiner模型进行了训练，则进行姿态的提炼预测，对结果进行评估计算dis，并且对dis反向传播
    if opt.refine_start:
        for ite in range(0, opt.iteration):
            pred_r, pred_t = refiner(new_points, emb, idx)
            dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_target, model_points, idx, new_points)
            dis.backward()

该篇博客，主要是对上面的refiner以及criterion_refine进行讲解，还是按顺序来，我们来看看refiner，假设这是第一执行refiner，那么其输入的new_points，emb来自于criterion函数的返回值。如果看了上篇博客，criterion返回的new_points，是经过逆转之后的points，如果看了前面的博客，可以知道这里的new_points相当于是points转换到相机的空间。其实这里可以理解为论文图示Figure 3的如下部分：

其中的points你可以理解（仅仅是理解，实际不是这样的）为下图：

中的transformed point cloud，然后呢，通过逆转，变成geometry embeddings = 代码中的new_points。现在呢我们拿到了geometry embeddings（new_points=空间信息），从图示中，我们可以知道，其还需要一个color embeddings（颜色信息），才能进行融合，如下：

这里的颜色信息从哪里来呢？他是由主干网络从当前特征抽取出来的color embeddings，也就是相当于源码中如下标记：

可以明确的看到，在迭代过程中emb一直都是从主干网络PoseNet抽取出来的color embeddings，没有做任何的改变，也符合了论文中的描述。但是这里的color embeddings也还是带有一些空间信息的，他的空间信息就是当前帧的空间信息。

前面说到new_points是points逆转之后的结果，但是color embeddings带有的少量空间信息是当前帧的（接近于points），所以出现了如下图示现象：

可以明显的看到，他们空间几何信息是不一样的。color embeddings=emb，geometry embeddings=new_points。他们既然空间信息不一样，那么我们就能对他进行pose的预测，对应论文图示如下部分：

可以看到color embeddings=emb（源码），geometry embeddings=new_points（源码）被送入了一个网络，然后进行姿态估算，也就是如下函数：

pred_r, pred_t = refiner(new_points, emb, idx)

得到新的估算姿态之后，他又和new_points（图示的current input point cloud）结合，把new_points进行偏转。本来new_points是points逆转过来的，又被偏转了回去，相当于又回到points状态，那么代码是怎么体现的呢？我们进入到criterion_refine函数，也就是lib/loss_refiner.py中的loss_calculation。我就不注释了啊 ，代码基本和lib/loss.py中的一样的，可以说是一个阉割的版本。如下代码

pred = torch.add(torch.bmm(model_points, base), pred_t)

就是一个转换的过程，对应图示如下部分

前面说到，转换过去之后，还需要一个逆转的过程，如下代码：

    ori_base = ori_base[0].view(1, 3, 3).contiguous()
    ori_t = t.repeat(bs * num_input_points, 1).contiguous().view(1, bs * num_input_points, 3)
    new_points = torch.bmm((points - ori_t), ori_base).contiguous()

    new_target = ori_target[0].view(1, num_point_mesh, 3).contiguous()
    ori_t = t.repeat(num_point_mesh, 1).contiguous().view(1, num_point_mesh, 3)
    new_target = torch.bmm((new_target - ori_t), ori_base).contiguous()

根据我们前面对代码的解析，看到- ori_t，我们就知道这是一个逆转的过程了。也就是图示如下部分：

这样，是不是带着大家走完一次迭代了。后续的迭代也是同样的原理。还又就是对其的理解了。上一篇博客中，其实已经举了一个很形象的例子。中的来说，就是说主干网络初始估算的pose，是存在误差的，这些误差的原因呢，要么就是语义分割做的不好（如切割下来的掺杂着背景），要么，存在遮挡现象等一系列的原因。为了进行减少这个误差，所以构建了Iterative Pose Refinement，通过多次迭代之后，用新估算出来的姿态，取弥补这些误差。

其实，这个地方比较抽象，不懂的朋友可以点个赞，然后来找我讨论，随时都欢迎。再见了，看到这里的你，应该也不容易，一路辛苦了。

6. 测试代码注释与详解

代码引导
在前面的章节，已经讲解了整个网络的数据预处理，迭代训练，网络框架，以及loss定义等，现在我们就剩下最后的网络测试没有讲解了，那么我们该章节就开始把，如果帅气貌美（我不知道是小哥哥还是小姐姐）的你已经看到了这里，要给我个赞啊，因为我们都一样好看嘛！

代码注释
通过前面的章节，我相信大家对整个框架流程已经十分透彻了，所以这里我直接给出测试代码注释tools/eval_linemod.py：

import _init_paths
import argparse
import os
import random
import numpy as np
import yaml
import copy
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from datasets.linemod.dataset import PoseDataset as PoseDataset_linemod
from lib.network import PoseNet, PoseRefineNet
from lib.loss import Loss
from lib.loss_refiner import Loss_refine
from lib.transformations import euler_matrix, quaternion_matrix, quaternion_from_matrix
from lib.knn.__init__ import KNearestNeighbor

parser = argparse.ArgumentParser()
# 验证数据根目录
parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir')
# 主干网络的模型，也就是PoseNet网络模型
parser.add_argument('--model', type=str, default = '',  help='resume PoseNet model')
# 姿态提炼网络模型，及PoseRefineNet网络模型
parser.add_argument('--refine_model', type=str, default = '',  help='resume PoseRefineNet model')
opt = parser.parse_args()

# 测试目标物体的总数目
num_objects = 13
objlist = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
# 根据当前需要测试帧的RGB-D，生成的点云数据，随机选择其中500个点云。
num_points = 500
# 提炼网络迭代的次数
iteration = 4
# batch_size
bs = 1
# 数据集配置网络的目录，该目录存放了models_info.yml文件，保存的x,y,轴的起始位置，以及其对应的半径（欧氏距离，作用后面有注释）
dataset_config_dir = 'datasets/linemod/dataset_config'
# 评估结果得输出目录
output_result_dir = 'experiments/eval_result/linemod'
# 和训练一样，主要为了寻找到那个最小得距离
knn = KNearestNeighbor(1)

# PoseNet网络模型得类构建，此时还没有进行前向传播
estimator = PoseNet(num_points = num_points, num_obj = num_objects)
estimator.cuda()

refiner = PoseRefineNet(num_points = num_points, num_obj = num_objects)
refiner.cuda()

# PoseNet模型参数加载-加载模型
estimator.load_state_dict(torch.load(opt.model))
# PoseRefineNet模型参数加载-加载模型
refiner.load_state_dict(torch.load(opt.refine_model))
estimator.eval()
refiner.eval()

# 以评估得模式加载数据
testdataset = PoseDataset_linemod('eval', num_points, False, opt.dataset_root, 0.0, True)
# 转化为torch迭代器形式
testdataloader = torch.utils.data.DataLoader(testdataset, batch_size=1, shuffle=False, num_workers=10)

# 对称物体索引的序列号
sym_list = testdataset.get_sym_list()
# 这里是获得测试数据中，点云最多的数目，实际为500
num_points_mesh = testdataset.get_num_points_mesh()
# 计算点云距离loss的类，此时只创建了一个类对象
criterion = Loss(num_points_mesh, sym_list)
criterion_refine = Loss_refine(num_points_mesh, sym_list)

# 存储每个模型的半径标准，该值主要和计算出来的dis（所有点云的平均距离）进行对比，如果小于该值，则数目达标了，否则该张图片的测试没有达标
diameter = []
meta_file = open('{0}/models_info.yml'.format(dataset_config_dir), 'r')
meta = yaml.load(meta_file)

# 除以1000应该是归一化，至于乘以0.1或许是单位转换（瞎猜的，大家可以去深究一下，知道答案了一定要告诉我啊）
for obj in objlist:
    diameter.append(meta[obj]['diameter'] / 1000.0 * 0.1)
print(diameter)

# 记录每个目标物体合格的数目，也就是测试出来的dis要小于models_info.yml文件对应的diameter参数，初始全部为0
success_count = [0 for i in range(num_objects)]
print(success_count)
# 该处是记录每个目标物体测试的总数目
num_count = [0 for i in range(num_objects)]

# 用于记录的日志
fw = open('{0}/eval_result_logs.txt'.format(output_result_dir), 'w')

for i, data in enumerate(testdataloader, 0):
    # 先获得一个样本需要的数据
    points, choose, img, target, model_points, idx = data
    if len(points.size()) == 2:
        print('No.{0} NOT Pass! Lost detection!'.format(i))
        fw.write('No.{0} NOT Pass! Lost detection!\n'.format(i))
        continue

    points, choose, img, target, model_points, idx = Variable(points).cuda(), \
                                                     Variable(choose).cuda(), \
                                                     Variable(img).cuda(), \
                                                     Variable(target).cuda(), \
                                                     Variable(model_points).cuda(), \
                                                     Variable(idx).cuda()

    # 通过PoseNet，预测当前帧的poses
    pred_r, pred_t, pred_c, emb = estimator(img, points, choose, idx)
    pred_r = pred_r / torch.norm(pred_r, dim=2).view(1, num_points, 1)
    pred_c = pred_c.view(bs, num_points)
    how_max, which_max = torch.max(pred_c, 1)
    pred_t = pred_t.view(bs * num_points, 1, 3)

    # 从所有的poses中找到最好的那一个，及置信度最高的那个
    my_r = pred_r[0][which_max[0]].view(-1).cpu().data.numpy()
    my_t = (points.view(bs * num_points, 1, 3) + pred_t)[which_max[0]].view(-1).cpu().data.numpy()
    my_pred = np.append(my_r, my_t)

    # 进行PoseRefineNet网络的循环迭代
    for ite in range(0, iteration):
        # 前面得到的结果已经不是torch格式的数据了，所以这里进行一个转换
        T = Variable(torch.from_numpy(my_t.astype(np.float32))).cuda().view(1, 3).repeat(num_points, 1).contiguous().view(1, num_points, 3)

        # 把姿态旋转参数r转换为矩阵形式
        my_mat = quaternion_matrix(my_r)
        R = Variable(torch.from_numpy(my_mat[:3, :3].astype(np.float32))).cuda().view(1, 3, 3)
        my_mat[0:3, 3] = my_t

        # 把points进行一个逆操作得到new_points（老铁们，不会为难我了吧，在前面的章节已经讲解很多次了）
        new_points = torch.bmm((points - T), R).contiguous()
        # 进行提炼，得到新的姿态pose
        pred_r, pred_t = refiner(new_points, emb, idx)


        pred_r = pred_r.view(1, 1, -1)
        pred_r = pred_r / (torch.norm(pred_r, dim=2).view(1, 1, 1))
        my_r_2 = pred_r.view(-1).cpu().data.numpy()
        my_t_2 = pred_t.view(-1).cpu().data.numpy()
        # 转化为四元数的矩阵
        my_mat_2 = quaternion_matrix(my_r_2)
        # 获得偏移参数（矩阵）
        my_mat_2[0:3, 3] = my_t_2

        # my_mat第一次迭代是主干网路初始的预测结果，之后是上一次迭代之后最后输出的结果
        # my_mat_2是当前迭代refiner预测出来的结果，矩阵相乘之后，把其当作该次迭代最后的结果
        my_mat_final = np.dot(my_mat, my_mat_2)
        my_r_final = copy.deepcopy(my_mat_final)
        my_r_final[0:3, 3] = 0

        # 同样转化为四元数矩阵
        my_r_final = quaternion_from_matrix(my_r_final, True)
        my_t_final = np.array([my_mat_final[0][3], my_mat_final[1][3], my_mat_final[2][3]])

        # 把当前最后的估算结果，赋值给送入下次迭代的pose
        my_pred = np.append(my_r_final, my_t_final)
        my_r = my_r_final
        my_t = my_t_final

    # Here 'my_pred' is the final pose estimation result after refinement ('my_r': quaternion, 'my_t': translation)
    model_points = model_points[0].cpu().detach().numpy()
    # 转化为四元数矩阵，取出其中的旋转矩阵，这里的my_r是最后预测的结果
    my_r = quaternion_matrix(my_r)[:3, :3]

    # model_points经过姿态转化之后的点云数据
    pred = np.dot(model_points, my_r.T) + my_t

    # 获得目标点云数据，其也是由model_points转换而来，但是其转换的pose是最标准的，
    target = target[0].cpu().detach().numpy()

    # 如果是对称的物体，需要做额外的处理，使用KNN找到最小的那个距离
    if idx[0].item() in sym_list:
        pred = torch.from_numpy(pred.astype(np.float32)).cuda().transpose(1, 0).contiguous()
        target = torch.from_numpy(target.astype(np.float32)).cuda().transpose(1, 0).contiguous()
        inds = knn(target.unsqueeze(0), pred.unsqueeze(0))
        target = torch.index_select(target, 1, inds.view(-1) - 1)
        dis = torch.mean(torch.norm((pred.transpose(1, 0) - target.transpose(1, 0)), dim=1), dim=0).item()
    else:
        dis = np.mean(np.linalg.norm(pred - target, axis=1))

    # 然后判断出该处理是否小于models_info.yml文件中保存的半径（欧氏距离）
    # 如果小于，表示就是合格的
    if dis < diameter[idx[0].item()]:
        success_count[idx[0].item()] += 1
        print('No.{0} Pass! Distance: {1}'.format(i, dis))
        fw.write('No.{0} Pass! Distance: {1}\n'.format(i, dis))
    else:
        print('No.{0} NOT Pass! Distance: {1}'.format(i, dis))
        fw.write('No.{0} NOT Pass! Distance: {1}\n'.format(i, dis))
    num_count[idx[0].item()] += 1

for i in range(num_objects):
    print('Object {0} success rate: {1}'.format(objlist[i], float(success_count[i]) / num_count[i]))
    fw.write('Object {0} success rate: {1}\n'.format(objlist[i], float(success_count[i]) / num_count[i]))
print('ALL success rate: {0}'.format(float(sum(success_count)) / sum(num_count)))
fw.write('ALL success rate: {0}\n'.format(float(sum(success_count)) / sum(num_count)))
fw.close()

本来想讲解一下的，我相信大家浏览一下就知道了，其实就是训练过程中的前向传播。