Update GaussianModel

scene/gaussian_model.py

@@ -721,12 +721,13 @@ def anchor_growing(self, iteration, grads, threshold, update_ratio, extra_ratio,
         for cur_level in range(self.levels):
             update_value = self.fork ** update_ratio
             level_mask = (self.get_level == cur_level).squeeze(dim=1)
+            level_ds_mask = (self.get_level == cur_level + 1).squeeze(dim=1)
             if torch.sum(level_mask) == 0:
                 continue
             cur_size = self.voxel_size / (float(self.fork) ** cur_level)
+            ds_size = cur_size / self.fork
             # update threshold
             cur_threshold = threshold * (update_value ** cur_level)
-            ds_size = cur_size / self.fork
             ds_threshold = cur_threshold * update_value
             extra_threshold = cur_threshold * extra_ratio
             # mask from grad threshold
@@ -752,7 +753,7 @@ def anchor_growing(self, iteration, grads, threshold, update_ratio, extra_ratio,
             selected_xyz = all_xyz.view([-1, 3])[candidate_mask]
             selected_grid_coords = torch.round((selected_xyz-self.init_pos)/cur_size).int()
             selected_grid_coords_unique, inverse_indices = torch.unique(selected_grid_coords, return_inverse=True, dim=0)
-            if selected_grid_coords_unique.shape[0] > 0:
+            if selected_grid_coords_unique.shape[0] > 0 and grid_coords.shape[0] > 0:
                 remove_duplicates = self.get_remove_duplicates(grid_coords, selected_grid_coords_unique)
                 remove_duplicates = ~remove_duplicates
                 candidate_anchor = selected_grid_coords_unique[remove_duplicates]*cur_size+self.init_pos
@@ -766,11 +767,11 @@ def anchor_growing(self, iteration, grads, threshold, update_ratio, extra_ratio,
                 new_level = torch.zeros([0], dtype=torch.int, device='cuda')
 
             if (~self.progressive or iteration > self.coarse_intervals[-1]) and cur_level < self.levels - 1:
-                grid_coords_ds = torch.round((self.get_anchor[level_mask]-self.init_pos)/ds_size).int()
+                grid_coords_ds = torch.round((self.get_anchor[level_ds_mask]-self.init_pos)/ds_size).int()
                 selected_xyz_ds = all_xyz.view([-1, 3])[candidate_ds_mask]
                 selected_grid_coords_ds = torch.round((selected_xyz_ds-self.init_pos)/ds_size).int()
                 selected_grid_coords_unique_ds, inverse_indices_ds = torch.unique(selected_grid_coords_ds, return_inverse=True, dim=0)
-                if selected_grid_coords_unique_ds.shape[0] > 0:
+                if selected_grid_coords_unique_ds.shape[0] > 0 and grid_coords_ds.shape[0] > 0:
                     remove_duplicates_ds = self.get_remove_duplicates(grid_coords_ds, selected_grid_coords_unique_ds)
                     remove_duplicates_ds = ~remove_duplicates_ds
                     candidate_anchor_ds = selected_grid_coords_unique_ds[remove_duplicates_ds]*ds_size+self.init_pos
@@ -848,7 +849,6 @@ def anchor_growing(self, iteration, grads, threshold, update_ratio, extra_ratio,
                 self._level = torch.cat([self._level, new_level], dim=0)
                 self._extra_level = torch.cat([self._extra_level, new_extra_level], dim=0)
 
-
     def adjust_anchor(self, iteration, check_interval=100, success_threshold=0.8, grad_threshold=0.0002, update_ratio=0.5, extra_ratio=4.0, extra_up=0.25, min_opacity=0.005):
         # # adding anchors
         grads = self.offset_gradient_accum / self.offset_denom # [N*k, 1]

single_train.sh

@@ -1,12 +1,12 @@
-scene="mipnerf360/garden"
+scene="db/playroom"
 exp_name="baseline"
 gpu=-1
 ratio=1
 resolution=-1
 appearance_dim=0
 
 fork=2
-base_layer=-1
+base_layer=12
 visible_threshold=0.9
 dist2level="round"
 update_ratio=0.2

train.py

@@ -26,7 +26,7 @@
 import wandb
 import time
 from os import makedirs
-import shutil, pathlib
+import shutil
 from pathlib import Path
 from PIL import Image
 import torchvision.transforms.functional as tf
@@ -70,8 +70,7 @@ def saveRuntimeCode(dst: str) -> None:
     for additionalPattern in additionalIgnorePatterns:
         ignorePatterns.append(additionalPattern)
 
-    log_dir = pathlib.Path(__file__).parent.resolve()
-
+    log_dir = Path(__file__).resolve().parent
 
     shutil.copytree(log_dir, dst, ignore=shutil.ignore_patterns(*ignorePatterns))
 

train.sh

@@ -39,27 +39,27 @@ time=$(date "+%Y-%m-%d_%H:%M:%S")
 
 if [ "$progressive" = "True" ]; then
     if [ "$warmup" = "True" ]; then
-        python train.py --eval -s data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} --warmup \
+        python train.py --eval -s /nas/shared/pjlab-lingjun-landmarks/renkerui/data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} --warmup \
         --iterations ${iterations} --port $port -m outputs/${data}/${logdir}/$time --appearance_dim ${appearance_dim} \
         --visible_threshold ${visible_threshold}  --base_layer ${base_layer} --dist2level ${dist2level} --update_ratio ${update_ratio} \
         --progressive --init_level ${init_level} --dist_ratio ${dist_ratio} --levels ${levels}  \
         --extra_ratio ${extra_ratio} --extra_up ${extra_up}
     else
-        python train.py --eval -s data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} \
+        python train.py --eval -s /nas/shared/pjlab-lingjun-landmarks/renkerui/data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} \
         --iterations ${iterations} --port $port -m outputs/${data}/${logdir}/$time --appearance_dim ${appearance_dim} \
         --visible_threshold ${visible_threshold}  --base_layer ${base_layer} --dist2level ${dist2level} --update_ratio ${update_ratio} \
         --progressive --init_level ${init_level} --dist_ratio ${dist_ratio} --levels ${levels}  \
         --extra_ratio ${extra_ratio} --extra_up ${extra_up}  
     fi
 else
     if [ "$warmup" = "True" ]; then
-        python train.py --eval -s data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} --warmup \
+        python train.py --eval -s /nas/shared/pjlab-lingjun-landmarks/renkerui/data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} --warmup \
         --iterations ${iterations} --port $port -m outputs/${data}/${logdir}/$time --appearance_dim ${appearance_dim} \
         --visible_threshold ${visible_threshold}  --base_layer ${base_layer} --dist2level ${dist2level} --update_ratio ${update_ratio} \
         --init_level ${init_level} --dist_ratio ${dist_ratio} --levels ${levels}  \
         --extra_ratio ${extra_ratio} --extra_up ${extra_up}  
     else
-        python train.py --eval -s data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} \
+        python train.py --eval -s /nas/shared/pjlab-lingjun-landmarks/renkerui/data/${data} -r ${resolution} --gpu ${gpu} --fork ${fork} --ratio ${ratio} \
         --iterations ${iterations} --port $port -m outputs/${data}/${logdir}/$time --appearance_dim ${appearance_dim} \
         --visible_threshold ${visible_threshold}  --base_layer ${base_layer} --dist2level ${dist2level} --update_ratio ${update_ratio} \
         --init_level ${init_level} --dist_ratio ${dist_ratio} --levels ${levels}  \

train_db.sh

@@ -6,7 +6,7 @@ appearance_dim=0
 
 fork=2
 base_layer=12
-visible_threshold=0.9 #0.9
+visible_threshold=0.9
 dist2level="round"
 update_ratio=0.2
 

train_mipnerf360.sh

@@ -5,7 +5,7 @@ resolution=-1
 appearance_dim=0
 
 fork=2
-base_layer=-1
+base_layer=10
 visible_threshold=-1 #0.9 
 dist2level="round"
 update_ratio=0.2