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计算。
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
