numObjects = 13
objList = [1, 2, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]
numPoints = 576
pooledImgSize = 48
batchSize = 1
datasetConfigDir = '/home/galen/deepLearning/poseEstimation/DenseFusion/datasets/linemod/Linemod_preprocessed/models/'
output_result_dir = 'experimentResult/eval_results/linemod'
knn = KNearestNeighbor(1)
estimator = poseNet(numPoints, numObjects)
estimator.cuda()
estimator.load_state_dict(torch.load(opt.model))
estimator.eval()
testDataset = PoseDataset('eval', numPoints, False, opt.datasetRoot, 0.0, True)
testDataLoader = torch.utils.data.DataLoader(testDataset,
batch_size=1,
shuffle=False,
num_workers=10)
symList = testDataset.get_sym_list()
numPointsMesh = testDataset.get_num_points_mesh()
poseNetLoss = Loss(numPointsMesh, symList)
diameter = []
meta_file = open('{0}/models_info.yml'.format(datasetConfigDir), 'r')
meta = yaml.load(meta_file)
kn = 0.1 # ADD 参数设置
for obj in objList:
diameter.append(meta[obj]['diameter'] / 1000.0 * kn)
def main():
print('------------')
opt.manualSeed = random.randint(1, 100) # 设定随机数
random.seed(opt.manualSeed) # 设定随机种子
torch.manual_seed(opt.manualSeed) # 设定随机种子
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 设定设备
opt.num_objects = 13 # 训练数据的物体种类数目
opt.num_points = 500 # 输入点云的数目
opt.outf = 'trained_models/linemod' # 训练模型保存的目录
opt.log_dir = 'logs/linemod' # log保存的目录
opt.repeat_epoch = 20 # 重复epoch数目
estimator = PoseNet(num_points=opt.num_points, num_obj=opt.num_objects) # 网络构建，构建完成，对物体的6D姿态进行预测
print(estimator)
estimator.to(device) # 选择设备
refiner = PoseRefineNet(num_points=opt.num_points, num_obj=opt.num_objects) # 对初步预测的姿态进行提炼
print(refiner)
refiner.to(device) # 选择设备
if opt.resume_posenet != '': # 对posenet模型的加载，如果有的话，
estimator.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_posenet)))
if opt.resume_refinenet != '': # 对refinenet模型的加载，如果有的话，
refiner.load_state_dict(torch.load('{0}/{1}'.format(opt.outf, opt.resume_refinenet)))
opt.refine_start = True # 标记refine网络开始训练
opt.decay_start = True # 标记refine网络参数开始衰减
opt.lr *= opt.lr_rate # 学习率变化
opt.w *= opt.w_rate # 权重衰减率变化
opt.batch_size = int(opt.batch_size / opt.iteration) # batchsize设定
optimizer = optim.Adam(refiner.parameters(), lr=opt.lr) # 设定refine的优化器
else:
opt.refine_start = False # 标记refine网络未开始训练
opt.decay_start = False # 标记refine网络参数未开始衰减
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr) # 设定posenet的优化器
# 加载训练数据集
dataset = PoseDataset('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)
# 加载验证数据集
test_dataset = PoseDataset('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
test_dataloder = 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) # loss计算
criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list) # refine_loss计算
best_test = np.Inf # 初始位置最好的模型,loss无限大
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))
# 记录每个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 # 用于计算平均距离
if opt.refine_start: # 如果refine模型，已经开始训练
estimator.eval()
refiner.train()
else:
estimator.train()
optimizer.zero_grad() # 优化器清零梯度
for rep in range(opt.repeat_epoch): # 每次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: 目标初始帧对应的点云信息
idx: 训练图片的下标
'''
# 将数据放到device上
points, choose, img, target, model_points, idx = points.to(device), choose.to(device), img.to(device), target.to(device), model_points.to(device), idx.to(device)
# 进行预测获得预测的姿态，和特征向量
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)
# 如果已经开始了refiner模型的训练
if opt.refine_start:
for iter in range(0, opt.iteration):
pred_r, pred_t = refiner(new_points, emb, idx) # 进行refiner预测
# 计算loss得到dis
dis, new_points, new_target = criterion_refine(pred_r, pred_t, new_target, model_points, idx, new_points)
dis.backward() # dis进行反向传播
else:
loss.backward() # 否则，则对loss进行反向传播
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.0
if train_count != 0 and train_count % 1000 == 0: # 模型保存
if opt.refine_start: # 已经开始refine模型
torch.save(refiner.state_dict(), '{0}/pose_refine_model_current.pth'.format(opt.outf))
else:
torch.save(estimator.state_dict(), '{0}.pos_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() # refiner模型
for j, data in enumerate(test_dataloder, 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: 目标初始帧对应的点云信息
idx: 训练图片的下标
'''
# 将数据放到device上
points, choose, img, target, model_points, idx = points.to(device), choose.to(device), img.to(
device), target.to(device), model_points.to(device), idx.to(device)
# 进行预测获得预测的姿态，和特征向量
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]
'''
# 对结果进行评估
_, 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 iter 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() # 用于计算平均距离
# 保存eval的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
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: # 保存refiner
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('----------------test model save finished-------------------')
# 参数变化
# 判断模型是否达到衰减要求
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)
# 模型没有达到loss阈值要求，refine_start = False，则修改相关参数，传递相关数，更新dataset和dataloader
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)
# 训练
dataset = PoseDataset('train', opt.num_points, True, opt.dataset_root, opt.noise_trans, opt.refine_start)
dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=opt.workers)
# 测试
test_dataset = PoseDataset('test', opt.num_points, False, opt.dataset_root, 0.0, opt.refine_start)
test_dataloder = 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)
iteration = 2
bs = 1
#dataset_config_dir = 'datasets/linemod/dataset_config' # for diameters 1.flansch 109 2. schaltgabel 153 3. stift 118
output_result_dir = 'experiments/eval_result/linemod'
knn = KNearestNeighbor(1)
estimator = PoseNet(num_points=num_points, num_obj=num_objects)
estimator.cuda()
refiner = PoseRefineNet(num_points=num_points, num_obj=num_objects)
refiner.cuda()
estimator.load_state_dict(torch.load(opt.model))
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)
testdataloader = torch.utils.data.DataLoader(testdataset,
batch_size=1,
shuffle=False,
num_workers=10)
sym_list = testdataset.get_sym_list()
num_points_mesh = testdataset.get_num_points_mesh()
criterion = Loss(num_points_mesh, sym_list)
criterion_refine = Loss_refine(num_points_mesh, sym_list)
diameter = [0.109, 0.153, 0.118]
success_count = [0 for i in range(num_objects)]
num_count = [0 for i in range(num_objects)]
fw = open('{0}/eval_result_logs.txt'.format(output_result_dir), 'w')
def main():
opt.manualSeed = random.randint(1, 10000)
random.seed(opt.manualSeed)
torch.manual_seed(
opt.manualSeed) #torch.manual_seed可以保证在种子值不变的情况下，每次torch.rand产生的结果相同
# torch.cuda.manual_seed(opt.manualSeed) torch.cuda.manual_seed_all(seed)
# 对于可重复的实验，有必要为任何使用随机数生成的进行随机种子设置。注意，cuDNN使用非确定性算法，并且可以使用torch.backends.cudnn.enabled = False来进行禁用。
# 根据数据集配置物体种类，用于预测位姿的点数，每个epoch训练的次数，模型保存的位置，log保存的位置
if opt.dataset == 'linemod':
opt.numObjects = 13
opt.numPoints = 576
opt.repeatEpoch = 20
opt.modelFolder = 'trainedModels/'
opt.logFolder = 'experimentResult/logs/'
else:
print('Unknown dataset')
return
# 定义网络，如果有训练过的模型，可加载； 确定网络优化方法
estimator = poseNet(opt.numPoints, opt.numObjects)
estimator.cuda() # or estimator.to('cuda')
if opt.resumePosenet != '':
estimator.load_state_dict(
torch.load('{0}/{1}'.format(opt.modelFolder, opt.resumePosenet)))
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
opt.decay_start = False
opt.refine_start = False
# 加载训练数据和测试数据
if opt.dataset == 'linemod':
dataset = PoseDataset('train', opt.numPoints, opt.addNoise,
opt.datasetRoot, opt.noiseTrans,
opt.refine_start)
test_dataset = PoseDataset('test', opt.numPoints, False,
opt.datasetRoot, 0.0, opt.refine_start)
dataLoader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=opt.workers)
testdataLoader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=opt.workers)
opt.symList = dataset.get_sym_list()
opt.numPointsMesh = 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.numPointsMesh,
opt.symList))
# 定义网络训练的误差
poseNetLoss = Loss(opt.numPointsMesh, opt.symList)
best_distance = np.Inf
if opt.startEpoch == 1:
for log in os.listdir(opt.logFolder):
os.remove(os.path.join(opt.logFolder, log))
st_time = time.time()
for epoch in range(opt.startEpoch, opt.nepoch):
# ------------------模型训练阶段------------------
logger = setup_logger(
'epoch%d' % epoch,
os.path.join(opt.logFolder, '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_sum = 0.0
estimator.train()
optimizer.zero_grad()
for rep in range(opt.repeatEpoch):
for i, data in enumerate(dataLoader, 0):
img, cloud, pointIndex, tarPoints, modelPoints, idx, ori_img = data
img = Variable(img).cuda()
cloud = Variable(cloud).cuda()
pointIndex = Variable(pointIndex).cuda()
tarPoints = Variable(tarPoints).cuda()
modelPoints = Variable(modelPoints).cuda()
idx = Variable(idx).cuda()
ori_img = np.array(ori_img)
pred_r, pred_t, pred_c, colorEmb = estimator(
img, cloud, pointIndex, idx)
loss, dis, newCloud, newTarPoints = poseNetLoss(
pred_r, pred_t, pred_c, tarPoints, modelPoints, idx, cloud,
opt.w, opt.refine_start)
loss.backward()
train_dis_sum += dis.item()
train_count += 1
if train_count % opt.batchSize == 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.batchSize),
train_count, train_dis_sum / opt.batchSize))
optimizer.step()
optimizer.zero_grad()
train_dis_sum = 0
if train_count != 0 and train_count % 1000 == 0:
torch.save(
estimator.state_dict(),
'{0}/pose_model_current.pth'.format(opt.modelFolder))
print(
'>>>>>>>>----------epoch {0} train finish---------<<<<<<<<'.format(
epoch))
# ------------------模型测试阶段------------------
logger = setup_logger(
'epoch%d_test' % epoch,
os.path.join(opt.logFolder, '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()
for j, data in enumerate(testdataLoader, 0):
img, cloud, pointIndex, tarPoints, modelPoints, idx, ori_img = data
img = Variable(img).cuda()
cloud = Variable(cloud).cuda()
pointIndex = Variable(pointIndex).cuda()
tarPoints = Variable(tarPoints).cuda()
modelPoints = Variable(modelPoints).cuda()
idx = Variable(idx).cuda()
ori_img = np.array(ori_img)
pred_r, pred_t, pred_c, colorEmb = estimator(
img, cloud, pointIndex, idx)
loss, dis, newCloud, newTarPoints = poseNetLoss(
pred_r, pred_t, pred_c, tarPoints, modelPoints, idx, cloud,
opt.w, opt.refine_start)
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_dis += dis.item()
test_count += 1
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_distance:
best_distance = test_dis
torch.save(
estimator.state_dict(),
'{0}/pose_model_{1}_{2}.pth'.format(opt.modelFolder, epoch,
test_dis))
print(epoch,
'>>>>>>>>----------BEST TEST MODEL SAVED---------<<<<<<<<')
if best_distance < opt.decayMargin and not opt.decay_start:
opt.decay_start = True
opt.lr *= opt.lr_dr
optimizer = optim.Adam(estimator.parameters(), lr=opt.lr)
def main():
opt.manualSeed = random.randint(1, 10000)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
opt.num_objects = 3
opt.num_points = 500
opt.outf = 'trained_models'
opt.log_dir = 'experiments/logs'
opt.repeat_epoch = 20
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()
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)
dataset = PoseDataset('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)
test_dataset = PoseDataset('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}'
.format(len(dataset), len(test_dataset), opt.num_points_mesh))
criterion = Loss(opt.num_points_mesh, opt.sym_list)
criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)
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):
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
if opt.refine_start:
estimator.eval() # affects dropout and batch normalization
refiner.train()
else:
estimator.train()
optimizer.zero_grad()
for rep in range(opt.repeat_epoch):
for i, data in enumerate(dataloader, 0):
points, choose, img, target, model_points, idx = data
#points ->torch.Size([500, 3]) ->在crop出来的像素区域随机选取500个点,利用相机内参结合深度值算出来的点云cloud
#choose ->torch.Size([1, 500])
#img ->torch.Size([3, 80, 80])
#target ->torch.Size([500, 3]) ->真实模型上随机选取的mesh点进行ground truth pose变换后得到的点
#model_points ->torch.Size([500, 3]) ->真实模型上随机选取的mesh点在进行pose变换前的点
#idx ->torch.Size([1])
#tensor([4], device='cuda:0')
#img和points对应rgb和点云信息,需要在网络内部fusion
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)
loss, dis, new_points, new_target = criterion(
pred_r, pred_t, pred_c, target, model_points, idx, points,
opt.w, opt.refine_start)
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
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):
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)
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()
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
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)
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)
dataset = PoseDataset('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)
test_dataset = PoseDataset('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}'
.format(len(dataset), len(test_dataset), opt.num_points_mesh))
criterion = Loss(opt.num_points_mesh, opt.sym_list)
criterion_refine = Loss_refine(opt.num_points_mesh, opt.sym_list)
estimator = PoseNet(num_points=num_points,
num_obj=num_objects) # PoseNet网络模型构建
estimator.to(device)
refiner = PoseRefineNet(num_points=num_points,
num_obj=num_objects) # PoseRefineNet网络模型构建
refiner.to(device)
estimator.load_state_dict(torch.load(opt.model)) # PoseNet模型参数加载
refiner.load_state_dict(torch.load(opt.refine_model)) # PoseRefineNet模型参数加载
estimator.eval()
refiner.eval()
# 以eval模式，加载数据
test_dataset = PoseDataset('eval', num_points, False, opt.dataset_root, 0.0,
True)
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False)
sym_list = test_dataset.get_sym_list() # 获取对称物体的索引
num_points_mesh = test_dataset.get_num_points_mesh() # 500
criterion = Loss(num_points_mesh, sym_list) # Loss加载
criterion_refine = Loss_refine(num_points_mesh, sym_list) # Loss_refine加载
diameter = [] # 存储模型给定的半径标准，与结果对比
meta_file = open(dataset_config_dir, 'r') # 读取model.info文件
meta = yaml.load(meta_file) # 加载model.info文件
for obj in objlist:
diameter.append(meta[obj]['diameter'] / 1000.0 * 0.1) # 存储到diameter中