def evaluate(self, epoch):
self.model.eval()
data_timer, model_timer = Timer(), Timer()
desc_loss_meter, det_loss_meter, acc_meter, d_pos_meter, d_neg_meter = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
num_iter = int(len(self.val_loader.dataset) // self.val_loader.batch_size)
num_iter = min(self.val_max_iter, num_iter)
test_loader_iter = self.val_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
inputs = test_loader_iter.next()
for k, v in inputs.items(): # load inputs to device.
if type(v) == list:
inputs[k] = [item.to(self.device) for item in v]
else:
inputs[k] = v.to(self.device)
data_timer.toc()
model_timer.tic()
features, scores = self.model(inputs)
anc_features = features[inputs["corr"][:, 0].long()]
pos_features = features[inputs["corr"][:, 1].long() + inputs['stack_lengths'][0][0]]
anc_scores = scores[inputs["corr"][:, 0].long()]
pos_scores = scores[inputs["corr"][:, 1].long() + inputs['stack_lengths'][0][0]]
desc_loss, acc, d_pos, d_neg, _, dist = self.evaluation_metric['desc_loss'](anc_features, pos_features, inputs['dist_keypts'])
det_loss = self.evaluation_metric['det_loss'](dist, anc_scores, pos_scores)
loss = desc_loss * self.metric_weight['desc_loss'] + det_loss * self.metric_weight['det_loss']
d_pos = np.mean(d_pos)
d_neg = np.mean(d_neg)
model_timer.toc()
desc_loss_meter.update(float(desc_loss))
det_loss_meter.update(float(det_loss))
d_pos_meter.update(float(d_pos))
d_neg_meter.update(float(d_neg))
acc_meter.update(float(acc))
if (iter + 1) % 100 == 0 and self.verbose:
print(f"Eval epoch: {epoch+1} [{iter+1:4d}/{num_iter}] "
f"desc loss: {desc_loss_meter.avg:.2f} "
f"det loss: {det_loss_meter.avg:.2f} "
f"acc: {acc_meter.avg:.2f} "
f"d_pos: {d_pos_meter.avg:.2f} "
f"d_neg: {d_neg_meter.avg:.2f} "
f"data time: {data_timer.avg:.2f}s "
f"model time: {model_timer.avg:.2f}s")
self.model.train()
res = {
'desc_loss': desc_loss_meter.avg,
'det_loss': det_loss_meter.avg,
'accuracy': acc_meter.avg,
'd_pos': d_pos_meter.avg,
'd_neg': d_neg_meter.avg,
}
print(f'Evaluation: Epoch {epoch}: Desc Loss {res["desc_loss"]}, Det Loss {res["det_loss"]}, Accuracy {res["accuracy"]}')
return res
def train_epoch(self, epoch):
data_timer, model_timer = Timer(), Timer()
loss_meter, acc_meter = AverageMeter(), AverageMeter()
num_iter = int(
len(self.train_loader.dataset) // self.train_loader.batch_size)
train_loader_iter = self.train_loader.__iter__()
# for iter, inputs in enumerate(self.train_loader):
for iter in range(num_iter):
data_timer.tic()
inputs = train_loader_iter.next()
for k, v in inputs.items(): # load inputs to device.
if type(v) == list:
inputs[k] = [item.to(self.device) for item in v]
else:
inputs[k] = v.to(self.device)
data_timer.toc()
model_timer.tic()
# forward
self.optimizer.zero_grad()
predict = self.model(inputs)
labels = inputs['labels'].long()
loss = self.evaluate_metric(predict, labels)
acc = torch.sum(
torch.max(predict, dim=1)[1].int() ==
labels.int()) * 100 / predict.shape[0]
# backward
loss.backward()
self.optimizer.step()
model_timer.toc()
loss_meter.update(float(loss))
acc_meter.update(float(acc))
if (iter + 1) % 1 == 0 and self.verbose:
print(f"Epoch: {epoch+1} [{iter+1:4d}/{num_iter}] "
f"loss: {loss_meter.avg:.2f} "
f"acc: {acc_meter.avg:.2f} "
f"data time: {data_timer.avg:.2f}s "
f"model time: {model_timer.avg:.2f}s")
# finish one epoch
epoch_time = model_timer.total_time + data_timer.total_time
self.train_hist['per_epoch_time'].append(epoch_time)
self.train_hist['loss'].append(loss_meter.avg)
self.train_hist['accuracy'].append(acc_meter.avg)
print(
f'Epoch {epoch+1}: Loss : {loss_meter.avg:.2f}, Accuracy: {acc_meter.avg:.2f} , Total time {epoch_time:.2f}s'
)
def evaluate(self, epoch):
self.model.eval()
data_timer, model_timer = Timer(), Timer()
loss_meter, acc_meter = AverageMeter(), AverageMeter()
num_iter = int(
len(self.test_loader.dataset) / self.test_loader.batch_size)
test_loader_iter = self.test_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
inputs = test_loader_iter.next()
for k, v in inputs.items(): # load inputs to device.
if type(v) == list:
inputs[k] = [item.to(self.device) for item in v]
else:
inputs[k] = v.to(self.device)
data_timer.toc()
model_timer.tic()
predict = self.model(inputs)
labels = inputs['labels'].long()
loss = self.evaluate_metric(predict, labels)
acc = torch.sum(
torch.max(predict, dim=1)[1].int() ==
labels.int()) * 100 / predict.shape[0]
model_timer.toc()
loss_meter.update(float(loss))
acc_meter.update(float(acc))
if (iter + 1) % 1 == 0 and self.verbose:
print(f"Eval epoch {epoch+1}: [{iter+1:3d}/{num_iter}] "
f"loss: {loss_meter.avg:.2f} "
f"acc: {acc_meter.avg:.2f} "
f"data time: {data_timer.avg:.2f}s "
f"model time: {model_timer.avg:.2f}s")
self.model.train()
res = {'loss': loss_meter.avg, 'accuracy': acc_meter.avg}
return res
def evaluate(self, epoch):
data_timer, model_timer = Timer(), Timer()
loss_meter, acc_meter, iou_meter = AverageMeter(), AverageMeter(
), AverageMeter()
num_iter = int(
len(self.test_loader.dataset) // self.train_loader.batch_size)
test_loader_iter = self.train_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
inputs = test_loader_iter.next()
for k, v in inputs.items(): # load inputs to device.
if type(v) == list:
inputs[k] = [item.to(self.device) for item in v]
else:
inputs[k] = v.to(self.device)
data_timer.toc()
model_timer.tic()
predict = self.model(inputs)
labels = inputs['labels'].long().squeeze()
loss = self.evaluate_metric(predict, labels)
# acc = torch.sum(torch.max(predict, dim=1)[1].int() == labels.int()) * 100 / predict.shape[0]
acc = calculate_acc(predict, labels)
part_iou = calculate_iou(predict,
labels,
stack_lengths=inputs['stack_lengths'][0],
n_parts=self.config.num_classes)
iou = np.mean(part_iou)
model_timer.toc()
loss_meter.update(float(loss))
iou_meter.update(float(iou))
acc_meter.update(float(acc))
if (iter + 1) % 1 == 0 and self.verbose:
print(f"Eval epoch: {epoch+1} [{iter+1:4d}/{num_iter}] "
f"loss: {loss_meter.avg:.2f} "
f"iou: {iou_meter.avg:.2f} "
f"data time: {data_timer.avg:.2f}s "
f"model time: {model_timer.avg:.2f}s")
self.model.train()
res = {
'iou': iou_meter.avg,
'loss': loss_meter.avg,
'accuracy': acc_meter.avg
}
return res
def train_epoch(self, epoch):
data_timer, model_timer = Timer(), Timer()
desc_loss_meter, det_loss_meter, acc_meter, d_pos_meter, d_neg_meter = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
num_iter = int(len(self.train_loader.dataset) // self.train_loader.batch_size)
num_iter = min(self.training_max_iter, num_iter)
train_loader_iter = self.train_loader.__iter__()
# for iter, inputs in enumerate(self.train_loader):
for iter in range(num_iter):
data_timer.tic()
inputs = train_loader_iter.next()
for k, v in inputs.items(): # load inputs to device.
if type(v) == list:
inputs[k] = [item.to(self.device) for item in v]
else:
inputs[k] = v.to(self.device)
data_timer.toc()
model_timer.tic()
# forward
self.optimizer.zero_grad()
features, scores = self.model(inputs)
anc_features = features[inputs["corr"][:, 0].long()]
pos_features = features[inputs["corr"][:, 1].long() + inputs['stack_lengths'][0][0]]
anc_scores = scores[inputs["corr"][:, 0].long()]
pos_scores = scores[inputs["corr"][:, 1].long() + inputs['stack_lengths'][0][0]]
desc_loss, acc, d_pos, d_neg, _, dist = self.evaluation_metric["desc_loss"](anc_features, pos_features, inputs['dist_keypts'])
det_loss = self.evaluation_metric['det_loss'](dist, anc_scores, pos_scores)
loss = desc_loss * self.metric_weight['desc_loss'] + det_loss * self.metric_weight['det_loss']
d_pos = np.mean(d_pos)
d_neg = np.mean(d_neg)
# backward
loss.backward()
do_step = True
for param in self.model.parameters():
if param.grad is not None:
if (1 - torch.isfinite(param.grad).long()).sum() > 0:
do_step = False
break
if do_step is True:
self.optimizer.step()
# if self.config.grad_clip_norm > 0:
# torch.nn.utils.clip_grad_value_(self.model.parameters(), self.config.grad_clip_norm)
model_timer.toc()
desc_loss_meter.update(float(desc_loss))
det_loss_meter.update(float(det_loss))
d_pos_meter.update(float(d_pos))
d_neg_meter.update(float(d_neg))
acc_meter.update(float(acc))
if (iter + 1) % 100 == 0 and self.verbose:
curr_iter = num_iter * (epoch - 1) + iter
self.writer.add_scalar('train/Desc_Loss', float(desc_loss_meter.avg), curr_iter)
self.writer.add_scalar('train/Det_Loss', float(det_loss_meter.avg), curr_iter)
self.writer.add_scalar('train/D_pos', float(d_pos_meter.avg), curr_iter)
self.writer.add_scalar('train/D_neg', float(d_neg_meter.avg), curr_iter)
self.writer.add_scalar('train/Accuracy', float(acc_meter.avg), curr_iter)
print(f"Epoch: {epoch} [{iter+1:4d}/{num_iter}] "
f"desc loss: {desc_loss_meter.avg:.2f} "
f"det loss: {det_loss_meter.avg:.2f} "
f"acc: {acc_meter.avg:.2f} "
f"d_pos: {d_pos_meter.avg:.2f} "
f"d_neg: {d_neg_meter.avg:.2f} "
f"data time: {data_timer.avg:.2f}s "
f"model time: {model_timer.avg:.2f}s")
# finish one epoch
epoch_time = model_timer.total_time + data_timer.total_time
print(f'Epoch {epoch}: Desc Loss: {desc_loss_meter.avg:.2f}, Det Loss : {det_loss_meter.avg:.2f}, Accuracy: {acc_meter.avg:.2f}, D_pos: {d_pos_meter.avg:.2f}, D_neg: {d_neg_meter.avg:.2f}, time {epoch_time:.2f}s')
def test_kitti(model, dataset, config):
# self.sess.run(dataset.test_init_op)
import sys
use_random_points = False
icp_save_path = "d3feat_output"
if use_random_points:
num_keypts = 5000
# icp_save_path = f'geometric_registration_kitti/D3Feat_{self.experiment_str}-rand{num_keypts}'
else:
num_keypts = 250
# icp_save_path = f'geometric_registration_kitti/D3Feat_{self.experiment_str}-pred{num_keypts}'
if not os.path.exists(icp_save_path):
os.mkdir(icp_save_path)
ch = logging.StreamHandler(sys.stdout)
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(format='%(asctime)s %(message)s',
datefmt='%m/%d %H:%M:%S',
handlers=[ch])
success_meter, loss_meter, rte_meter, rre_meter = AverageMeter(
), AverageMeter(), AverageMeter(), AverageMeter()
feat_timer, reg_timer = Timer(), Timer()
for i in range(dataset.length):
import pdb
pdb.set_trace()
# feat_timer.tic()
ops = [
model.anchor_inputs, model.out_features, model.out_scores,
model.anc_id, model.pos_id, model.accuracy
]
[inputs, features, scores, anc_id, pos_id,
accuracy] = self.sess.run(ops, {model.dropout_prob: 1.0})
# feat_timer.toc()
# print(accuracy, anc_id)
stack_lengths = inputs['stack_lengths']
first_pcd_indices = np.arange(stack_lengths[0])
second_pcd_indices = np.arange(stack_lengths[1]) + stack_lengths[0]
# anc_points = inputs['points'][0][first_pcd_indices]
# pos_points = inputs['points'][0][second_pcd_indices]
# anc_features = features[first_pcd_indices]
# pos_features = features[second_pcd_indices]
# anc_scores = scores[first_pcd_indices]
# pos_scores = scores[second_pcd_indices]
if use_random_points:
anc_keypoints_id = np.random.choice(stack_lengths[0],
num_keypts)
pos_keypoints_id = np.random.choice(
stack_lengths[1], num_keypts) + stack_lengths[0]
anc_points = inputs['points'][0][anc_keypoints_id]
pos_points = inputs['points'][0][pos_keypoints_id]
anc_features = features[anc_keypoints_id]
pos_features = features[pos_keypoints_id]
anc_scores = scores[anc_keypoints_id]
pos_scores = scores[pos_keypoints_id]
else:
scores_anc_pcd = scores[first_pcd_indices]
scores_pos_pcd = scores[second_pcd_indices]
anc_keypoints_id = np.argsort(scores_anc_pcd,
axis=0)[-num_keypts:].squeeze()
pos_keypoints_id = np.argsort(
scores_pos_pcd,
axis=0)[-num_keypts:].squeeze() + stack_lengths[0]
anc_points = inputs['points'][0][anc_keypoints_id]
anc_features = features[anc_keypoints_id]
anc_scores = scores[anc_keypoints_id]
pos_points = inputs['points'][0][pos_keypoints_id]
pos_features = features[pos_keypoints_id]
pos_scores = scores[pos_keypoints_id]
pcd0 = make_open3d_point_cloud(anc_points)
pcd1 = make_open3d_point_cloud(pos_points)
feat0 = make_open3d_feature(anc_features, 32,
anc_features.shape[0])
feat1 = make_open3d_feature(pos_features, 32,
pos_features.shape[0])
reg_timer.tic()
filename = anc_id.decode("utf-8") + "-" + pos_id.decode(
"utf-8").split("@")[-1] + '.npz'
if os.path.exists(join(icp_save_path, filename)):
data = np.load(join(icp_save_path, filename))
T_ransac = data['trans']
print(f"Read from {join(icp_save_path, filename)}")
else:
distance_threshold = dataset.voxel_size * 1.0
ransac_result = open3d.registration.registration_ransac_based_on_feature_matching(
pcd0, pcd1, feat0, feat1, distance_threshold,
open3d.registration.TransformationEstimationPointToPoint(
False), 4, [
open3d.registration.
CorrespondenceCheckerBasedOnEdgeLength(0.9),
open3d.registration.
CorrespondenceCheckerBasedOnDistance(
distance_threshold)
],
open3d.registration.RANSACConvergenceCriteria(50000, 1000)
# open3d.registration.RANSACConvergenceCriteria(4000000, 10000)
)
# print(ransac_result)
T_ransac = ransac_result.transformation.astype(np.float32)
np.savez(join(icp_save_path, filename),
trans=T_ransac,
anc_pts=anc_points,
pos_pts=pos_points,
anc_scores=anc_scores,
pos_scores=pos_scores)
reg_timer.toc()
T_gth = inputs['trans']
# loss_ransac = corr_dist(T_ransac, T_gth, anc_points, pos_points, weight=None, max_dist=1)
loss_ransac = 0
rte = np.linalg.norm(T_ransac[:3, 3] - T_gth[:3, 3])
rre = np.arccos(
(np.trace(T_ransac[:3, :3].transpose() @ T_gth[:3, :3]) - 1) /
2)
if rte < 2:
rte_meter.update(rte)
if not np.isnan(rre) and rre < np.pi / 180 * 5:
rre_meter.update(rre * 180 / np.pi)
if rte < 2 and not np.isnan(rre) and rre < np.pi / 180 * 5:
success_meter.update(1)
else:
success_meter.update(0)
logging.info(
f"{anc_id} Failed with RTE: {rte}, RRE: {rre * 180 / np.pi}"
)
loss_meter.update(loss_ransac)
if (i + 1) % 10 == 0:
logging.info(
f"{i+1} / {dataset.num_test}: Feat time: {feat_timer.avg},"
+
f" Reg time: {reg_timer.avg}, Loss: {loss_meter.avg}, RTE: {rte_meter.avg},"
+
f" RRE: {rre_meter.avg}, Success: {success_meter.sum} / {success_meter.count}"
+ f" ({success_meter.avg * 100} %)")
feat_timer.reset()
reg_timer.reset()
logging.info(
f"Total loss: {loss_meter.avg}, RTE: {rte_meter.avg}, var: {rte_meter.var},"
+
f" RRE: {rre_meter.avg}, var: {rre_meter.var}, Success: {success_meter.sum} "
+ f"/ {success_meter.count} ({success_meter.avg * 100} %)")
def register_one_scene(inlier_ratio_threshold, distance_threshold, save_path,
return_dict, scene):
gt_matches = 0
pred_matches = 0
keyptspath = f"{save_path}/keypoints/{scene}"
descpath = f"{save_path}/descriptors/{scene}"
scorepath = f"{save_path}/scores/{scene}"
gtpath = f'geometric_registration/gt_result/{scene}-evaluation/'
gtLog = loadlog(gtpath)
inlier_num_meter, inlier_ratio_meter = AverageMeter(), AverageMeter()
pcdpath = f"{config.root}/fragments/{scene}/"
num_frag = len([
filename for filename in os.listdir(pcdpath)
if filename.endswith('ply')
])
for id1 in range(num_frag):
for id2 in range(id1 + 1, num_frag):
cloud_bin_s = f'cloud_bin_{id1}'
cloud_bin_t = f'cloud_bin_{id2}'
key = f"{id1}_{id2}"
if key not in gtLog.keys():
# skip the pairs that have less than 30% overlap.
num_inliers = 0
inlier_ratio = 0
gt_flag = 0
else:
source_keypts = get_keypts(keyptspath, cloud_bin_s)
target_keypts = get_keypts(keyptspath, cloud_bin_t)
source_desc = get_desc(descpath, cloud_bin_s, 'D3Feat')
target_desc = get_desc(descpath, cloud_bin_t, 'D3Feat')
source_score = get_scores(scorepath, cloud_bin_s,
'D3Feat').squeeze()
target_score = get_scores(scorepath, cloud_bin_t,
'D3Feat').squeeze()
source_desc = np.nan_to_num(source_desc)
target_desc = np.nan_to_num(target_desc)
# randomly select 5000 keypts
if args.random_points:
source_indices = np.random.choice(
range(source_keypts.shape[0]), args.num_points)
target_indices = np.random.choice(
range(target_keypts.shape[0]), args.num_points)
else:
source_indices = np.argsort(
source_score)[-args.num_points:]
target_indices = np.argsort(
target_score)[-args.num_points:]
source_keypts = source_keypts[source_indices, :]
source_desc = source_desc[source_indices, :]
target_keypts = target_keypts[target_indices, :]
target_desc = target_desc[target_indices, :]
corr = build_correspondence(source_desc, target_desc)
gt_trans = gtLog[key]
frag1 = source_keypts[corr[:, 0]]
frag2_pc = o3d.geometry.PointCloud()
frag2_pc.points = o3d.utility.Vector3dVector(
target_keypts[corr[:, 1]])
frag2_pc.transform(gt_trans)
frag2 = np.asarray(frag2_pc.points)
distance = np.sqrt(np.sum(np.power(frag1 - frag2, 2), axis=1))
num_inliers = np.sum(distance < distance_threshold)
inlier_ratio = num_inliers / len(distance)
if inlier_ratio > inlier_ratio_threshold:
pred_matches += 1
gt_matches += 1
inlier_num_meter.update(num_inliers)
inlier_ratio_meter.update(inlier_ratio)
recall = pred_matches * 100.0 / gt_matches
return_dict[scene] = [recall, inlier_num_meter.avg, inlier_ratio_meter.avg]
logging.info(
f"{scene}: Recall={recall:.2f}%, inlier ratio={inlier_ratio_meter.avg*100:.2f}%, inlier num={inlier_num_meter.avg:.2f}"
)
return recall, inlier_num_meter.avg, inlier_ratio_meter.avg
def train_epoch(self, epoch):
# create meters and timers
meter_list = ['class_loss', 'trans_loss', 'sm_loss', 'reg_recall', 're', 'te', 'precision', 'recall', 'f1']
meter_dict = {}
for key in meter_list:
meter_dict[key] = AverageMeter()
data_timer, model_timer = Timer(), Timer()
num_iter = int(len(self.train_loader.dataset) / self.batch_size)
num_iter = min(self.training_max_iter, num_iter)
trainer_loader_iter = self.train_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
(corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels) = trainer_loader_iter.next()
if self.gpu_mode:
corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels = \
corr_pos.cuda(), src_keypts.cuda(), tgt_keypts.cuda(), gt_trans.cuda(), gt_labels.cuda()
data = {
'corr_pos': corr_pos,
'src_keypts': src_keypts,
'tgt_keypts': tgt_keypts,
}
data_timer.toc()
model_timer.tic()
# forward
self.optimizer.zero_grad()
res = self.model(data)
pred_trans, pred_labels = res['final_trans'], res['final_labels']
# classification loss
class_stats = self.evaluate_metric['ClassificationLoss'](pred_labels, gt_labels)
class_loss = class_stats['loss']
# spectral matching loss
sm_loss = self.evaluate_metric['SpectralMatchingLoss'](res['M'], gt_labels)
# transformation loss
trans_loss, reg_recall, re, te, rmse = self.evaluate_metric['TransformationLoss'](pred_trans, gt_trans, src_keypts, tgt_keypts, pred_labels)
loss = self.metric_weight['ClassificationLoss'] * class_loss + self.metric_weight['SpectralMatchingLoss'] * sm_loss
if epoch > self.transformation_loss_start_epoch and self.metric_weight['TransformationLoss'] > 0.0:
loss += self.metric_weight['TransformationLoss'] * trans_loss
stats = {
'class_loss': float(class_loss),
'sm_loss': float(sm_loss),
'trans_loss': float(trans_loss),
'reg_recall': float(reg_recall),
're': float(re),
'te': float(te),
'precision': class_stats['precision'],
'recall': class_stats['recall'],
'f1': class_stats['f1'],
}
# backward
loss.backward()
do_step = True
for param in self.model.parameters():
if param.grad is not None:
if (1 - torch.isfinite(param.grad).long()).sum() > 0:
do_step = False
break
if do_step is True:
self.optimizer.step()
model_timer.toc()
if not np.isnan(float(loss)):
for key in meter_list:
if not np.isnan(stats[key]):
meter_dict[key].update(stats[key])
else: # debug the loss calculation process.
import pdb
pdb.set_trace()
if (iter + 1) % 100 == 0 and self.verbose:
curr_iter = num_iter * (epoch - 1) + iter
for key in meter_list:
self.writer.add_scalar(f"Train/{key}", meter_dict[key].avg, curr_iter)
print(f"Epoch: {epoch} [{iter+1:4d}/{num_iter}] "
f"sm_loss: {meter_dict['sm_loss'].avg:.2f} "
f"class_loss: {meter_dict['class_loss'].avg:.2f} "
f"trans_loss: {meter_dict['trans_loss'].avg:.2f} "
f"reg_recall: {meter_dict['reg_recall'].avg:.2f}% "
f"re: {meter_dict['re'].avg:.2f}degree "
f"te: {meter_dict['te'].avg:.2f}cm "
f"data_time: {data_timer.avg:.2f}s "
f"model_time: {model_timer.avg:.2f}s "
)
def evaluate(self, epoch):
self.model.eval()
# create meters and timers
meter_list = ['class_loss', 'trans_loss', 'sm_loss', 'reg_recall', 're', 'te', 'precision', 'recall', 'f1']
meter_dict = {}
for key in meter_list:
meter_dict[key] = AverageMeter()
data_timer, model_timer = Timer(), Timer()
num_iter = int(len(self.val_loader.dataset) / self.batch_size)
num_iter = min(self.val_max_iter, num_iter)
val_loader_iter = self.val_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
(corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels) = val_loader_iter.next()
if self.gpu_mode:
corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels = \
corr_pos.cuda(), src_keypts.cuda(), tgt_keypts.cuda(), gt_trans.cuda(), gt_labels.cuda()
data = {
'corr_pos': corr_pos,
'src_keypts': src_keypts,
'tgt_keypts': tgt_keypts,
}
data_timer.toc()
model_timer.tic()
# forward
res = self.model(data)
pred_trans, pred_labels = res['final_trans'], res['final_labels']
# classification loss
class_stats = self.evaluate_metric['ClassificationLoss'](pred_labels, gt_labels)
class_loss = class_stats['loss']
# spectral matching loss
sm_loss = self.evaluate_metric['SpectralMatchingLoss'](res['M'], gt_labels)
# transformation loss
trans_loss, reg_recall, re, te, rmse = self.evaluate_metric['TransformationLoss'](pred_trans, gt_trans, src_keypts, tgt_keypts, pred_labels)
model_timer.toc()
stats = {
'class_loss': float(class_loss),
'sm_loss': float(sm_loss),
'trans_loss': float(trans_loss),
'reg_recall': float(reg_recall),
're': float(re),
'te': float(re),
'precision': class_stats['precision'],
'recall': class_stats['recall'],
'f1': class_stats['f1'],
}
for key in meter_list:
if not np.isnan(stats[key]):
meter_dict[key].update(stats[key])
self.model.train()
res = {
'sm_loss': meter_dict['sm_loss'].avg,
'class_loss': meter_dict['class_loss'].avg,
'reg_recall': meter_dict['reg_recall'].avg,
'trans_loss': meter_dict['trans_loss'].avg,
}
for key in meter_list:
self.writer.add_scalar(f"Val/{key}", meter_dict[key].avg, epoch)
return res
def train_epoch(self, epoch):
# create meters and timers
meter_list = [
'class_loss', 'trans_loss', 'sm_loss', 'reg_recall', 're', 'te',
'precision', 'recall', 'f1'
]
meter_dict = {}
for key in meter_list:
meter_dict[key] = AverageMeter()
data_timer, model_timer = Timer(), Timer()
num_iter = len(self.train_loader)
num_iter = min(self.training_max_iter, num_iter)
self.train_loader.sampler.set_epoch(epoch)
trainer_loader_iter = self.train_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
input_dict = trainer_loader_iter.next()
(corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels, _,
_) = self.train_feature_extractor.process_batch(input_dict)
if self.gpu_mode:
corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels = \
corr_pos.cuda(), src_keypts.cuda(), tgt_keypts.cuda(), gt_trans.cuda(), gt_labels.cuda()
data = {
'corr_pos': corr_pos,
'src_keypts': src_keypts,
'tgt_keypts': tgt_keypts,
}
data_timer.toc()
model_timer.tic()
# forward
self.optimizer.zero_grad()
res = self.model(data)
pred_trans, pred_labels = res['final_trans'], res['final_labels']
# classification loss
class_stats = self.evaluate_metric['ClassificationLoss'](
pred_labels, gt_labels)
class_loss = class_stats['loss']
# spectral matching loss
sm_loss = self.evaluate_metric['SpectralMatchingLoss'](res['M'],
gt_labels)
# transformation loss
trans_loss, reg_recall, re, te, rmse = self.evaluate_metric[
'TransformationLoss'](pred_trans, gt_trans, src_keypts,
tgt_keypts, pred_labels)
loss = self.metric_weight[
'ClassificationLoss'] * class_loss + self.metric_weight[
'SpectralMatchingLoss'] * sm_loss
if epoch > self.transformation_loss_start_epoch and self.metric_weight[
'TransformationLoss'] > 0.0:
loss += self.metric_weight['TransformationLoss'] * trans_loss
stats = {
'class_loss': float(class_loss),
'sm_loss': float(sm_loss),
'trans_loss': float(trans_loss),
'reg_recall': float(reg_recall),
're': float(re),
'te': float(te),
'precision': class_stats['precision'],
'recall': class_stats['recall'],
'f1': class_stats['f1'],
}
# backward
loss.backward()
self.sum_gradients()
do_step = True
for param in self.model.parameters():
if param.grad is not None:
if (1 - torch.isfinite(param.grad).long()).sum() > 0:
do_step = False
break
if do_step is True:
self.optimizer.step()
model_timer.toc()
if not np.isnan(float(loss)):
for key in meter_list:
if not np.isnan(stats[key]):
meter_dict[key].update(stats[key])
else: # debug the loss calculation process.
print("Bug found, ignoring.")
# import pdb
# pdb.set_trace()
if (iter + 1) % 50 == 0 and self.verbose:
curr_iter = num_iter * (epoch - 1) + iter
report = torch.tensor([
1.0, data_timer.avg, model_timer.avg,
meter_dict['class_loss'].avg, meter_dict['trans_loss'].avg,
meter_dict['sm_loss'].avg, meter_dict['reg_recall'].avg,
meter_dict['re'].avg, meter_dict['te'].avg,
meter_dict['precision'].avg, meter_dict['recall'].avg,
meter_dict['f1'].avg
],
device=torch.cuda.current_device())
dist.all_reduce(report, op=dist.ReduceOp.SUM)
if self.rank == 0:
mean_meter_dict = {}
count = report[0].item()
data_timer_avg = report[1].item() / count
model_timer_avg = report[2].item() / count
mean_meter_dict['class_loss'] = report[3].item() / count
mean_meter_dict['trans_loss'] = report[4].item() / count
mean_meter_dict['sm_loss'] = report[5].item() / count
mean_meter_dict['reg_recall'] = report[6].item() / count
mean_meter_dict['re'] = report[7].item() / count
mean_meter_dict['te'] = report[8].item() / count
mean_meter_dict['precision'] = report[9].item() / count
mean_meter_dict['recall'] = report[10].item() / count
mean_meter_dict['f1'] = report[11].item() / count
for key in meter_list:
self.writer.add_scalar(f"Train/{key}",
mean_meter_dict[key], curr_iter)
print(
f"{time.strftime('%m/%d %H:%M:%S')} Epoch: {epoch} [{iter+1:4d}/{num_iter}] "
f"sm_loss: {mean_meter_dict['sm_loss']:.2f} "
f"class_loss: {mean_meter_dict['class_loss']:.2f} "
f"trans_loss: {mean_meter_dict['trans_loss']:.2f} "
f"reg_recall: {mean_meter_dict['reg_recall']:.2f}% "
f"re: {mean_meter_dict['re']:.2f}degree "
f"te: {mean_meter_dict['te']:.2f}cm "
f"data_time: {data_timer_avg:.2f}s "
f"model_time: {model_timer_avg:.2f}s ")
def evaluate(self, epoch):
self.model.eval()
# create meters and timers
meter_list = [
'class_loss', 'trans_loss', 'sm_loss', 'reg_recall', 're', 'te',
'precision', 'recall', 'f1'
]
meter_dict = {}
for key in meter_list:
meter_dict[key] = AverageMeter()
data_timer, model_timer = Timer(), Timer()
num_iter = len(self.val_loader)
num_iter = min(self.val_max_iter, num_iter)
val_loader_iter = self.val_loader.__iter__()
for iter in range(num_iter):
data_timer.tic()
input_dict = val_loader_iter.next()
(corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels, _,
_) = self.val_feature_extractor.process_batch(input_dict)
if self.gpu_mode:
corr_pos, src_keypts, tgt_keypts, gt_trans, gt_labels = \
corr_pos.cuda(), src_keypts.cuda(), tgt_keypts.cuda(), gt_trans.cuda(), gt_labels.cuda()
data = {
'corr_pos': corr_pos,
'src_keypts': src_keypts,
'tgt_keypts': tgt_keypts,
}
data_timer.toc()
model_timer.tic()
# forward
res = self.model(data)
pred_trans, pred_labels = res['final_trans'], res['final_labels']
# classification loss
class_stats = self.evaluate_metric['ClassificationLoss'](
pred_labels, gt_labels)
class_loss = class_stats['loss']
# spectral matching loss
sm_loss = self.evaluate_metric['SpectralMatchingLoss'](res['M'],
gt_labels)
# transformation loss
trans_loss, reg_recall, re, te, rmse = self.evaluate_metric[
'TransformationLoss'](pred_trans, gt_trans, src_keypts,
tgt_keypts, pred_labels)
model_timer.toc()
stats = {
'class_loss': float(class_loss),
'sm_loss': float(sm_loss),
'trans_loss': float(trans_loss),
'reg_recall': float(reg_recall),
're': float(re),
'te': float(te),
'precision': class_stats['precision'],
'recall': class_stats['recall'],
'f1': class_stats['f1'],
}
for key in meter_list:
if not np.isnan(stats[key]):
meter_dict[key].update(stats[key])
self.model.train()
report = torch.tensor([
1.0, meter_dict['class_loss'].avg, meter_dict['trans_loss'].avg,
meter_dict['sm_loss'].avg, meter_dict['reg_recall'].avg,
meter_dict['re'].avg, meter_dict['te'].avg,
meter_dict['precision'].avg, meter_dict['recall'].avg,
meter_dict['f1'].avg
],
device=torch.cuda.current_device())
dist.all_reduce(report, op=dist.ReduceOp.SUM)
mean_meter_dict = {}
count = report[0].item()
mean_meter_dict['class_loss'] = report[1].item() / count
mean_meter_dict['trans_loss'] = report[2].item() / count
mean_meter_dict['sm_loss'] = report[3].item() / count
mean_meter_dict['reg_recall'] = report[4].item() / count
mean_meter_dict['re'] = report[5].item() / count
mean_meter_dict['te'] = report[6].item() / count
mean_meter_dict['precision'] = report[7].item() / count
mean_meter_dict['recall'] = report[8].item() / count
mean_meter_dict['f1'] = report[9].item() / count
res = {
'sm_loss': mean_meter_dict['sm_loss'],
'class_loss': mean_meter_dict['class_loss'],
'reg_recall': mean_meter_dict['reg_recall'],
'trans_loss': mean_meter_dict['trans_loss'],
}
if self.rank == 0:
for key in meter_list:
self.writer.add_scalar(f"Val/{key}", mean_meter_dict[key],
epoch)
return res