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

batchlize transfomer and add forward_train/test in pipeline

Browse Source
This commit is contained in:
hofee
2024-08-22 00:41:30 +08:00
parent b06dede4b8
commit 8f82bec66c
2 changed files with 78 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
core/pipeline.py
View File

@@ -1,4 +1,4 @@
import torch
from torch import nn
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
@@ -24,8 +24,43 @@ 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
random_t = random_t.unsqueeze(-1)
mu, std = self.view_finder.marginal_prob(gt_delta_rot_6d, random_t)
std = std.view(-1, 1)
z = torch.randn_like(gt_delta_rot_6d)
perturbed_x = mu + z * std
target_score = - z * std / (std ** 2)
return perturbed_x, random_t, target_score, std
def forward_train(self, data):
output = {}
pts_list = data['pts_list']
pose_list = data['pose_list']
gt_delta_rot_6d = data["delta_rot_6d"]
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)
''' get std '''
perturbed_x, random_t, target_score, std = self.pertube_data(gt_delta_rot_6d)
input_data = {
"sampled_pose": perturbed_x,
"t": random_t,
"seq_feat": seq_feat,
}
estimated_score = self.view_finder(input_data)
output = {
"estimated_score": estimated_score,
"target_score": target_score,
"std": std
}
return output
def forward_test(self,data):
pts_list = data['pts_list']
pose_list = data['pose_list']
pts_feat_list = []
@@ -34,9 +69,10 @@ class NBVReconstructionPipeline(nn.Module):
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)
output['estimated_score'] = self.view_finder.next_best_view(seq_feat)
estimated_delta_rot_6d, in_process_sample = self.view_finder.next_best_view(seq_feat)
result = {
"estimated_delta_rot_6d": estimated_delta_rot_6d,
"in_process_sample": in_process_sample
}
return result
return output
def forward_test(self,data):
pass