diff --git a/configs/local/inference_config.yaml b/configs/local/inference_config.yaml index 384e912..bfc30f0 100644 --- a/configs/local/inference_config.yaml +++ b/configs/local/inference_config.yaml @@ -62,7 +62,7 @@ pipeline: module: pointnet++_encoder: in_dim: 3 - params_name: dense + params_name: light pointnet_encoder: in_dim: 3 diff --git a/configs/server/server_train_config.yaml b/configs/server/server_train_config.yaml index 87ba6e3..79ea4a1 100644 --- a/configs/server/server_train_config.yaml +++ b/configs/server/server_train_config.yaml @@ -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 diff --git a/core/pipeline.py b/core/pipeline.py index ae04d9e..1050628 100644 --- a/core/pipeline.py +++ b/core/pipeline.py @@ -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 diff --git a/modules/pointnet++_encoder.py b/modules/pointnet++_encoder.py index c597fb5..ee8cc3c 100644 --- a/modules/pointnet++_encoder.py +++ b/modules/pointnet++_encoder.py @@ -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)