hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix bug for training

Browse Source
This commit is contained in:
hofee
2024-09-12 15:11:09 +08:00
parent a79ca7749d
commit 4c69ed777b
15 changed files with 201 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
core/pipeline.py
View File

@@ -14,12 +14,11 @@ class NBVReconstructionPipeline(nn.Module):
self.pose_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["pose_encoder"])
self.seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["seq_encoder"])
self.view_finder = ComponentFactory.create(namespace.Stereotype.MODULE, config["view_finder"])
self.eps = 1e-5
def forward(self, data):
mode = data["mode"]
# ----- Debug Trace ----- #
import ipdb; ipdb.set_trace()
# ------------------------ #
if mode == namespace.Mode.TRAIN:
return self.forward_train(data)
elif mode == namespace.Mode.TEST:
@@ -27,29 +26,22 @@ class NBVReconstructionPipeline(nn.Module):
else:
Log.error("Unknown mode: {}".format(mode), True)
def pertube_data(self, gt_delta_rot_6d):
bs = gt_delta_rot_6d.shape[0]
random_t = torch.rand(bs, device=self.device) * (1. - self.eps) + self.eps
def pertube_data(self, gt_delta_9d):
bs = gt_delta_9d.shape[0]
random_t = torch.rand(bs, device=gt_delta_9d.device) * (1. - self.eps) + self.eps
random_t = random_t.unsqueeze(-1)
mu, std = self.view_finder.marginal_prob(gt_delta_rot_6d, random_t)
mu, std = self.view_finder.marginal_prob(gt_delta_9d, random_t)
std = std.view(-1, 1)
z = torch.randn_like(gt_delta_rot_6d)
z = torch.randn_like(gt_delta_9d)
perturbed_x = mu + z * std
target_score = - z * std / (std ** 2)
return perturbed_x, random_t, target_score, std
def forward_train(self, data):
pts_list = data['pts_list']
pose_list = data['pose_list']
gt_rot_6d = data["nbv_cam_pose"]
pts_feat_list = []
pose_feat_list = []
for pts,pose in zip(pts_list,pose_list):
pts_feat_list.append(self.pts_encoder.encode_points(pts))
pose_feat_list.append(self.pose_encoder.encode_pose(pose))
seq_feat = self.seq_encoder.encode_sequence(pts_feat_list, pose_feat_list)
seq_feat = self.get_seq_feat(data)
''' get std '''
perturbed_x, random_t, target_score, std = self.pertube_data(gt_rot_6d)
best_to_1_pose_9d_batch = data["best_to_1_pose_9d"]
perturbed_x, random_t, target_score, std = self.pertube_data(best_to_1_pose_9d_batch)
input_data = {
"sampled_pose": perturbed_x,
"t": random_t,
@@ -64,14 +56,7 @@ class NBVReconstructionPipeline(nn.Module):
return output
def forward_test(self,data):
pts_list = data['pts_list']
pose_list = data['pose_list']
pts_feat_list = []
pose_feat_list = []
for pts,pose in zip(pts_list,pose_list):
pts_feat_list.append(self.pts_encoder.encode_points(pts))
pose_feat_list.append(self.pose_encoder.encode_pose(pose))
seq_feat = self.seq_encoder.encode_sequence(pts_feat_list, pose_feat_list)
seq_feat = self.get_seq_feat(data)
estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(seq_feat)
result = {
"pred_pose_9d": estimated_delta_rot_9d,
@@ -79,4 +64,19 @@ class NBVReconstructionPipeline(nn.Module):
}
return result
def get_seq_feat(self, data):
scanned_pts_batch = data['scanned_pts']
scanned_n_to_1_pose_9d_batch = data['scanned_n_to_1_pose_9d']
best_to_1_pose_9d_batch = data["best_to_1_pose_9d"]
pts_feat_seq_list = []
pose_feat_seq_list = []
for scanned_pts,scanned_n_to_1_pose_9d in zip(scanned_pts_batch,scanned_n_to_1_pose_9d_batch):
print(scanned_n_to_1_pose_9d.shape)
scanned_pts = scanned_pts.to(best_to_1_pose_9d_batch.device)
scanned_n_to_1_pose_9d = scanned_n_to_1_pose_9d.to(best_to_1_pose_9d_batch.device)
pts_feat_seq_list.append(self.pts_encoder.encode_points(scanned_pts))
pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_1_pose_9d))
seq_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
return seq_feat