misc

Author: JasonLSC

examples/simple_trainer.py

@@ -58,17 +58,12 @@ def __init__(self):
             torch.profiler.ProfilerActivity.CUDA,
         ]
 
-        # 基础配置
-        self.wait = 1      # 开始记录前等待的步数
-        self.warmup = 2    # 预热步数
-        self.active = 30_000    # 实际分析的步数
-        # self.repeat = 2    # 重复次数
-        # self.skip_first = 10  # 跳过前N步（可选）
+        self.wait = 1 
+        self.warmup = 2 
+        self.active = 30_000 
 
-        # 创建schedule
         self.schedule = self._create_schedule()
 
-        # 其他profiler设置
         self.on_trace_ready = torch.profiler.tensorboard_trace_handler('./log/profiler')
         self.record_shapes = True
         self.profile_memory = True
@@ -79,8 +74,6 @@ def _create_schedule(self):
             wait=self.wait,
             warmup=self.warmup,
             active=self.active,
-            # repeat=self.repeat,
-            # skip_first=self.skip_first
         )
 
     def update_schedule(self, **kwargs):

gsplat/compression/outlier_filter.py

@@ -6,32 +6,28 @@
 from torch import Tensor
 
 def filter_splats(splats: Dict[str, Tensor], opa_thres: float=0.005, std_factor: float=2.0, k_neighbors: int=10):
-    # 1. 首先基于不透明度过滤
     opacity_mask = torch.sigmoid(splats["opacities"]) >= opa_thres
 
-    # # 2. 使用KD树计算每个点到其k个最近邻的平均距离
+    
     # pos = splats["means"].cpu().numpy()
     # kdtree = KDTree(pos)
     # distances, _ = kdtree.query(pos, k=k_neighbors)
     # mean_distances = np.mean(distances, axis=1)
 
-    # # 3. 计算平均距离的统计特征
+    #
     # dist_mean = np.mean(mean_distances)
     # dist_std = np.std(mean_distances)
 
-    # # 4. 基于距离判定离群点
     # distance_mask = mean_distances <= (dist_mean + std_factor * dist_std)
     # distance_mask = torch.from_numpy(distance_mask).to(opacity_mask.device)
 
-    # # 5. 组合两个过滤条件
-    # ## v1: 保守方案
+    # ## v1: 
     # # outlier = torch.logical_and(~opacity_mask, ~distance_mask) # 既在位置上离群，且不透明度又低
     # # valid_mask = ~outlier
-    # ## v2: 激进方案
+    # ## v2:
     # # valid_mask = torch.logical_and(opacity_mask, distance_mask)
     valid_mask = opacity_mask
 
-    # 6. 逐一过滤元素
     for n, v in splats.items():
         splats[n] = v[valid_mask]
 

gsplat/compression_simulation/ada_mask.py

@@ -20,25 +20,21 @@ def get_temperature(self, current_step):
         if current_step < self.annealing_start_iter:
             return self.start_temp
 
-        # 计算退火后的温度
         progress = (current_step - self.annealing_start_iter) / \
                   (self.total_iters - self.annealing_start_iter)
-        progress = min(max(progress, 0), 1)  # 限制在[0,1]范围内
+        progress = min(max(progress, 0), 1) 
 
-        # 使用指数衰减
         temperature = self.start_temp * math.exp(
             math.log(self.end_temp / self.start_temp) * progress
         )
         return temperature
 
     def forward(self, x, current_step):
         if self.training:
-            # 训练时使用当前温度的sigmoid
             self.current_iter = current_step
             temperature = self.get_temperature(current_step)
             mask = torch.sigmoid(self.mask_logits / temperature)
         else:
-            # 测试时使用硬阈值
             mask = (torch.sigmoid(self.mask_logits) >= 0.5).float()
         return x * mask
 
@@ -47,21 +43,13 @@ def get_binary_mask(self):
         return (torch.sigmoid(self.mask_logits) >= 0.5).float()
 
     def get_sparsity_loss(self, lambda_l1=0.01, lambda_target=0.1):
-        """
-        计算稀疏性损失，包括：
-        1. L1正则化
-        2. 目标稀疏度的KL散度损失
-        """
         temperature = self.get_temperature(self.current_iter)
         mask = torch.sigmoid(self.mask_logits / temperature)
 
-        # L1 正则化
         l1_loss = lambda_l1 * torch.mean(mask)
 
-        # 计算当前的平均激活率
         current_sparsity = torch.mean(mask)
 
-        # KL散度损失，确保整体稀疏度接近目标值
         target = torch.tensor(self.target_sparsity).to(mask.device)
         kl_loss = lambda_target * F.binary_cross_entropy(current_sparsity, target)
 

gsplat/compression_simulation/entropy_model.py

@@ -161,10 +161,9 @@ def forward(self, x, Q=None, **kwargs):
         else:
             Q = torch.tensor([Q], device=x.device)
 
-        # [N, C] -> [C, 1, N], batch维度移到最后以提高内存访问效率
+        # [N, C] -> [C, 1, N]
         x = x.t().unsqueeze(1)
 
-        # 预计算公共部分
         half_Q = 0.5 * Q.detach()
         x_lower = x - half_Q
         x_upper = x + half_Q
@@ -209,10 +208,9 @@ def forward(self, x, Q=None, **kwargs):
         else:
             Q = torch.tensor([Q], device=x.device)
 
-        # [N, C] -> [C, 1, N], batch维度移到最后以提高内存访问效率
+        # [N, C] -> [C, 1, N]
         x = x.t().unsqueeze(1)
 
-        # 预计算公共部分
         half_Q = 0.5 * Q.detach()
         x_lower = x - half_Q
         x_upper = x + half_Q
@@ -265,10 +263,9 @@ def get_likelihood(self, x, Q=None, **kwargs):
         else:
             Q = torch.tensor([Q], device=x.device)
 
-        # [N, C] -> [C, 1, N], batch维度移到最后以提高内存访问效率
+        # [N, C] -> [C, 1, N], 
         x = x.t().unsqueeze(1)
 
-        # 预计算公共部分
         half_Q = 0.5 * Q.detach()
         x_lower = x - half_Q
         x_upper = x + half_Q
@@ -346,23 +343,6 @@ def get_means_and_scales(self, pos):
         scales = torch.clamp(scales, min=1e-9)
 
         return means, scales
-
-# class Low_bound(torch.autograd.Function):
-#     @staticmethod
-#     def forward(ctx, x):
-#         ctx.save_for_backward(x)
-#         x = torch.clamp(x, min=1e-6)
-#         return x
-
-#     @staticmethod
-#     def backward(ctx, g):
-#         x, = ctx.saved_tensors
-#         grad1 = g.clone()
-#         grad1[x < 1e-6] = 0
-#         pass_through_if = np.logical_or(
-#             x.cpu().numpy() >= 1e-6, g.cpu().numpy() < 0.0)
-#         t = torch.Tensor(pass_through_if+0.0).cuda()
-#         return grad1 * t
 
 def lower_bound_fwd(x: Tensor, bound: Tensor) -> Tensor:
     return torch.max(x, bound)