update strong p++

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

configs/server/server_train_config.yaml

@@ -7,13 +7,13 @@ 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:
@@ -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

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

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()
+    net = PointNet2Encoder(config={"in_dim": 3, "params_name": "strong"}).cuda()
-    pts = torch.randn(2, 1024, 3).cuda()
+    pts = torch.randn(2, 2444, 3).cuda()
     print(torch.mean(pts, dim=1))
     pre = net.encode_points(pts)
     print(pre.shape)