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

update strong p++

Browse Source
This commit is contained in:
hofee 2025-01-23 08:58:10 +00:00
parent 5a03659112
commit d7fb64ed13
4 changed files with 37 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
configs/local/inference_config.yaml
View File

@ -62,7 +62,7 @@ pipeline:
module:
pointnet++_encoder:
in_dim: 3
params_name: dense
params_name: light
pointnet_encoder:
in_dim: 3

12
configs/server/server_train_config.yaml
View File

@ -7,19 +7,19 @@ runner:
parallel: False
experiment:
name: train_ab_global_only_with_wp_p++_dense
name: train_ab_global_only_with_wp_p++_strong
root_dir: "experiments"
use_checkpoint: False
epoch: -1 # -1 stands for last epoch
max_epochs: 5000
save_checkpoint_interval: 1
test_first: True
test_first: False
train:
optimizer:
type: Adam
lr: 0.0001
losses:
losses:
- gf_loss
dataset: OmniObject3d_train
test:
@ -39,7 +39,7 @@ dataset:
type: train
cache: True
ratio: 1
batch_size: 80
batch_size: 64
num_workers: 128
pts_num: 8192
load_from_preprocess: True
@ -98,7 +98,7 @@ module:
pointnet++_encoder:
in_dim: 3
params_name: dense
params_name: strong
transformer_seq_encoder:
embed_dim: 256
@ -110,7 +110,7 @@ module:
gf_view_finder:
t_feat_dim: 128
pose_feat_dim: 256
main_feat_dim: 2048
main_feat_dim: 5120
regression_head: Rx_Ry_and_T
pose_mode: rot_matrix
per_point_feature: False

2
core/pipeline.py
View File

@ -75,7 +75,7 @@ class NBVReconstructionPipeline(nn.Module):
def forward_test(self, data):
main_feat = self.get_main_feat(data)
repeat_num = data.get("repeat_num", 100)
repeat_num = data.get("repeat_num", 1)
main_feat = main_feat.repeat(repeat_num, 1)
estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(
main_feat

31
modules/pointnet++_encoder.py
View File

@ -33,6 +33,29 @@ ClsMSG_CFG_Light = {
'DP_RATIO': 0.5,
}
ClsMSG_CFG_Light_2048 = {
'NPOINTS': [512, 256, 128, None],
'RADIUS': [[0.02, 0.04], [0.04, 0.08], [0.08, 0.16], [None, None]],
'NSAMPLE': [[16, 32], [16, 32], [16, 32], [None, None]],
'MLPS': [[[16, 16, 32], [32, 32, 64]],
[[64, 64, 128], [64, 96, 128]],
[[128, 196, 256], [128, 196, 256]],
[[256, 256, 1024], [256, 512, 1024]]],
'DP_RATIO': 0.5,
}
ClsMSG_CFG_Strong = {
'NPOINTS': [1024, 512, 256, 128, None], # 增加采样点，获取更多细节
'RADIUS': [[0.02, 0.05], [0.05, 0.1], [0.1, 0.2], [0.2, 0.4], [None, None]], # 增大感受野
'NSAMPLE': [[32, 64], [32, 64], [32, 64], [32, 64], [None, None]], # 提高每层的采样点数
'MLPS': [[[32, 32, 64], [64, 64, 128]], # 增强 MLP 层，增加特征提取能力
[[128, 128, 256], [128, 128, 256]],
[[256, 256, 512], [256, 384, 512]],
[[512, 512, 1024], [512, 768, 1024]],
[[1024, 1024, 2048], [1024, 1024, 2048]]], # 增加更深的特征层
'DP_RATIO': 0.4, # Dropout 比率稍微降低，以保留更多信息
}
ClsMSG_CFG_Lighter = {
'NPOINTS': [512, 256, 128, 64, None],
'RADIUS': [[0.01], [0.02], [0.04], [0.08], [None]],
@ -53,6 +76,10 @@ def select_params(name):
return ClsMSG_CFG_Lighter
elif name == 'dense':
return ClsMSG_CFG_Dense
elif name == 'light_2048':
return ClsMSG_CFG_Light_2048
elif name == 'strong':
return ClsMSG_CFG_Strong
else:
raise NotImplementedError
@ -114,8 +141,8 @@ if __name__ == '__main__':
seed = 100
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
net = PointNet2Encoder(config={"in_dim": 3, "params_name": "light"}).cuda()
pts = torch.randn(2, 1024, 3).cuda()
net = PointNet2Encoder(config={"in_dim": 3, "params_name": "strong"}).cuda()
pts = torch.randn(2, 2444, 3).cuda()
print(torch.mean(pts, dim=1))
pre = net.encode_points(pts)
print(pre.shape)