【DenseFusion代码详解】训练过程train.py

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

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

---

train过程包括train和test，在每一次迭代之后都对本轮迭代的模型进行test。代码位置tools/train.py

首先是一些超参数的设置：

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default = 'ycb', help='ycb or linemod')
parser.add_argument('--dataset_root', type=str, default = '', help='dataset root dir (''YCB_Video_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')
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')
parser.add_argument('--iteration', type=int, default = 2, help='number of refinement iterations')
parser.add_argument('--nepoch', type=int, default=500, help='max number of epochs to train')
parser.add_argument('--resume_posenet', type=str, default = '',  help='resume PoseNet model')
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()

argparse 模块是一种命令行接口，也就是在.sh文件或者命令行输入相应的参数，赋予opt中相应变量的具体值。上述变量分别的含义如下：

--dataset：数据集，选择YCB或者LineMOD，默认为YCB

--dataset_root：数据集的路径

--batch_size：批量大小，但这里训练的时候都是1，可能是因为每次分割出的图片大小都不一样，只能一个一个训练

--workers：读取数据的进程数量，PyTorch的 DataLoader 允许使用多进程来加速数据读取

--lr：学习率

--lr_rate：学习率衰减率

--w：平衡超参数

--w_rate：权重衰减率

--decay_margin：衰减阈值

--refine_margin：开始迭代自优化refine的阈值

--noise_trans：添加到训练数据中随机噪声的范围

--iteration：迭代自优化的次数

--nepoch：最大训练周期

--resume_posenet：之前训练已经保存的posenet模型

--resume_refinenet：之前训练已经保存的refinenet模型

--start_epoch：开始训练的epoch

下面是main()函数中的内容。

首先设置随机数种子，用于参数初始化。

def main():
    opt.manualSeed = random.randint(1, 10000)
    random.seed(opt.manualSeed)
    torch.manual_seed(opt.manualSeed)

下面定义数据集，包括物体类别数，随机选取的点云数、保存路径等：

    #选择数据集
    if opt.dataset == 'ycb':
        opt.num_objects = 21 #数据集中物体的类别数
        opt.num_points = 1000 #随机筛选点云的点数
        opt.outf = 'trained_models/ycb' #保存训练模型的路径
        opt.log_dir = 'experiments/logs/ycb' #保存log文件的路径
        opt.repeat_epoch = 1 #number of repeat times for one epoch training
    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

这里区别两个数据集，物体类别数分别为21和13，输入点云的点数分别为1000和500，分别将model和log保存在各自的文件夹下面。

    #选择网络
    estimator = PoseNet(num_points = opt.num_points, num_obj = opt.num_objects)
    estimator.cuda()
    refiner = PoseRefineNet(num_points = opt.num_points, num_obj = opt.num_objects)
    refiner.cuda()

estimator为PoseNet网络，即用于预测姿态的主干网络，refiner为PoseRefineNet网络，用于后续迭代自优化。

    #是否加载前面训练的posenet模型
    if opt.resume_posenet != '':
        estimator.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet)))
    #是否加载前面训练的refinenet模型
    if opt.resume_refinenet != '':
        refiner.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet)))
        opt.refine_start = True #开始refine过程
        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 #还没开始refine过程
        opt.decay_start = False #还没开始衰减
        optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)

opt.outf是各自数据集保存的模型路径，上述代码是因为前面的训练过程中可能会发生中断，但会保存训练的模型，在下次训练中如果使用--resume_posenet指定先前训练的posenet模型路径，就会加载先前训练的模型继续训练，--resume_refinenet也一样，如果指定先前训练的模型地址，则会加载模型，并设置学习率和权重衰减，否则视为还没开始refine和衰减过程。

    #加载训练数据集
    if opt.dataset == 'ycb':
        dataset = PoseDataset_ycb('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 == '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)

加载训练数据集和测试数据集，Dataloader使用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))

获取该数据集的对称物体编号列表、mesh点数。 

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

对loss进行初始化，分别定义loss和loss_refine，详见loss.py和loss_refiner.py。

    best_test = np.Inf

将最好模型的loss值best_test设置成无穷大。

    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()

如果开始训练的epoch为1，则视为重头开始训练，就将之前训练的log文件全都删除。并记录开始时间。

然后进入epoch循环，下面的所有代码都在这个循环里面。

    for epoch in range(opt.start_epoch, opt.nepoch): #开始训练的epoch和最大的epoch
        #保存每次训练的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() #将梯度初始化为0

如果开始refine过程了，那么姿势估计网络posenet开始eval模式，迭代自优化网络开始训练train，否则，posenet还是训练模式train。

        #每个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, img, target, model_points, idx = Variable(points).cuda(), \
                                                                 Variable(choose).cuda(), \
                                                                 Variable(img).cuda(), \
                                                                 Variable(target).cuda(), \
                                                                 Variable(model_points).cuda(), \
                                                                 Variable(idx).cuda()

enumerate() 函数用于将一个可遍历的数据对象(如列表、元组或字符串)组合为一个索引序列，同时列出数据和数据下标，这个时候会调用dataset的__getitem__函数，获取一项数据，详见dataset.py。这里获取的变量含义如下：

points：由深度图转换成点云并随机筛选500个点，相机坐标系。

choose：所选择500个点云的索引，[bs, 1, 500]

img：通过语义分割之后剪切下来的RGB图像

target：根据model_points点云信息，以及标准旋转偏移矩阵转换过的目标点云[bs,500,3]

model_points：目标初始帧（模型）对应的点云信息[bs,500,3]

idx：目标物体类别

获取这些数据之后，将截取的RGB图像、筛选的点云、索引和物体类别输入到PoseNet姿态估计网络中进行训练，详见network.py。

                #使用PoseNet进行姿势估计
                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]，置信度，每个像素都有一个预测

emb: 经过choose操作之后的img，与点云一一对应

 得到预测的pose和置信度之后，开始计算loss：

                #计算loss
                loss, dis, new_points, new_target = criterion(pred_r, pred_t, pred_c, target, model_points, idx, points, opt.w, opt.refine_start)

将预测值、目标点云、初始帧点云模型、编号、筛选的500个点云、平衡超参数等作为输入计算loss，详见loss.py。

                if opt.refine_start: #如果开始了refine过程
                    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()

这个部分是refine的部分，如果没有开始refine就直接对loss进行反向传播，如果开始了refine过程，则将上述loss计算输出的由预测pose和points逆转而来的new_points作为PoseRefineNet网络的输入，与经过choose之后的rbg图像一起，进行网络的训练，具体训练过程详见network.py。网络输出纠正的pose，然后计算refine过程的loss，这个loss只有一个像素的输出，该像素置信度最大，同样进行反向传播。而这个refine过程可以设置循环次数iteration，默认为2。

                train_dis_avg += dis.item()
                train_count += 1

                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))

上述代码实现每一个batch输出log信息，每1000次训练保存一个模型，如果已有refine过程则保存refine模型，如果没有则保存estimator模型。

到此训练过程结束。下面开始测试过程。

        #保存每次测试的log文件
        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()

训练完train样本后，生成的模型model要用来测试样本。这里须将模型设置为eval模式，否则的话，有输入数据，即使不训练，它也会改变权值。这是model中含有BN层和Dropout所带来的的性质。具体参考Pytorch基础 | eval()的用法比较_公众号机器学习与生成对抗网络的博客-CSDN博客

        #下面是对测试数据集进行测试的过程
        for j, data in enumerate(testdataloader, 0):
            #获取测试数据的各个值
            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()

 同样地，对测试数据进行预处理，获取用于测试的点云、RGB等，进行格式转换。

            #进行PoseNet姿势估计
            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)

用PoseNet计算姿态，对姿态计算loss（但这里只输出dis，为最大置信度像素的loss）。 

            #如果开始了refine过程
            if opt.refine_start:
                for ite in range(0, opt.iteration):
                    #计算每次refine的pose和loss
                    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)

如果有refine，则将上一次预测姿态逆转的点云作为输入，用PoseRefineNet计算新的pose，然后计算refine过程的loss（也就是dis） 

            #输出log
            test_dis += dis.item()
            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，输出log
        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))

将所有dis相加求平均，得到测试过程的平均dis，测试过程结束，到此，就完成了每次epoch的训练和测试步骤，下面的代码是为了下次epoch做准备。

        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---------<<<<<<<<')

如果测试的dis小于最好的dis（初始best_dis为无穷大），就将test_dis作为best_dis，然后保存本次epoch最好的模型（如果有refine过程就保存refiner，没有就保存estimator）。

下面的代码是用来判断是否开始权重衰减，一旦开始了之后，就不会再执行以下的代码。

        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过程，和衰减过程一样，一旦开始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)
            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 == '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)

可以看作，refine过程的迭代是在全局一次迭代之中进行的，但这里的batch_size对加载数据并没有什么影响，因为所有DataLoader中的batch_size都为1，可能是因为语义分割之后的图片大小都不一样，只能一个一个处理。然后重新加载refine过程的数据，定义refine过程的loss计算。