refactor: implement adaptive masking strategies in compression simulation

Author: JasonLSC

dev_docs/config_usage.md

@@ -35,3 +35,57 @@ python examples/simple_trainer.py default --config logs/run_001.yaml
 
 > 提示：策略（`strategy`）字段会以 `{type: DefaultStrategy, params: {...}}` 形式写入 YAML，方便直接调整或复用。
 
+
+## Adaptive Mask 重构设计（进行中）
+
+- **目标**：将自适应 SHN mask 重构为独立子模块，和可微分量化、熵约束保持同等抽象层级；Trainer 与 orchestrator 仅通过统一接口消费 mask 输出/损失/指标。
+- **配置结构**：
+  ```python
+  @dataclass
+  class LearnableMaskSettings:
+      start_temp: float = 5.0
+      end_temp: float = 0.1
+      total_iters: int = 30_000
+      target_sparsity: float = 0.2
+      lr: float = 1e-2
+
+  @dataclass
+  class GradientMaskSettings:
+      grad_threshold: float = 2e-3  # 小于该阈值的梯度会被清零
+
+  @dataclass
+  class MaskConfig:
+      enabled: bool = False
+      strategy: Optional[str] = "learnable"
+      start_step: int = 10_000
+      learnable: LearnableMaskSettings = field(default_factory=LearnableMaskSettings)
+      gradient: GradientMaskSettings = field(default_factory=GradientMaskSettings)
+  ```
+  YAML 示例：
+  ```yaml
+  compression_sim_cfg:
+    mask:
+      enabled: true
+      strategy: gradient
+      start_step: 8000
+      gradient:
+        grad_threshold: 0.003
+  ```
+- **模块划分**：新增 `gsplat/compression_simulation/mask.py`，定义 `MaskResult`、`AdaptiveMaskBase` 接口；`AdaptiveMaskFactory` 根据策略返回 `LearnableAdaptiveMask`（封装 `AnnealingMask`）、`GradientAdaptiveMask`（梯度阈值裁剪）或 `NullAdaptiveMask`。
+- **Gradient 策略**：
+  - `maybe_update` 统计 SHN 非零比例，供日志与阈值计算。
+  - `apply` 在 `step > start_step` 时注册一次 `tensor.register_hook`，返回 `MaskResult(value=tensor, metrics={"mask_ratio": ..., "mask_grad_threshold": ...})`。
+  - Hook 逻辑：
+    ```python
+    def grad_hook(grad):
+        shn = tensor.detach()
+        zero_mask = (shn == 0).all(dim=-1).all(dim=-1)
+        grad_norm = grad.flatten(2).norm(p=2, dim=-1)
+        mask = ~(zero_mask & (grad_norm < cfg.gradient.grad_threshold))
+        while mask.ndim < grad.ndim:
+            mask = mask.unsqueeze(-1)
+        return grad * mask
+    ```
+- **Orchestrator 接入**：`CompressionSimulation.run()` 中调用 `mask.maybe_update(step, splats)`；对 `shN` 应用 `mask.apply` 并汇总 `loss_terms` 与 `metrics`；`step_optimizers`/`state_dict` 同步 mask 状态。
+- **Trainer/日志**：训练循环只消费 `SimulationResult.loss_terms.get("mask")` 与 `metrics["mask_ratio"]`；导出 PLY 时通过 mask 模块暴露的 `get_binary_mask()` 获取最终掩码。
+- **测试计划**：新增 `tests/test_adaptive_mask.py` 覆盖 learnable（温度调度、loss、优化器）、gradient（阈值裁剪、指标）与 null（passthrough）。

dev_docs/post-training_comp.md

@@ -0,0 +1,132 @@
+# Post-Training Compression 重构设想
+
+## 目标
+- 支持统一的后处理压缩入口，可消费训练得到的 splats（来源为 PLY 或 checkpoint）。
+- 提供 encode / decode 两个核心函数：
+  - `encode()`：从输入载入 splats，执行 pruning、quantization、3D→2D 映射、视频/图像编码，输出码流或压缩文件集合。
+  - `decode()`：从码流恢复 splats，执行视频/图像解码、反量化等操作，可输出内存结构或保存为 PLY。
+- 允许通过配置选择 codec 以及预处理策略，便于扩展新变种。
+
+## 拟议架构
+### 1. 顶层 Orchestrator（暂名 `PostTrainingCompressor`）
+- 初始化参数：
+  - `input_spec`: 指定输入类型（`ply` 或 `ckpt`）及路径。
+  - `codec_config`: 指定编码方式（PNG / Entropy / HEVC / SeqYUV 等）及对应参数。
+  - `preprocess_config`: 可选，控制 pruning、排序、属性变换等。
+  - `quant_config`: 可选，控制按属性的量化位宽、截断范围等。
+- 方法：
+  - `encode()`：驱动全流程；返回压缩输出路径、元信息。
+  - `decode(compressed_dir)`：读取码流、还原 splats；可返回字典或写入 PLY。
+
+### 2. 数据加载层
+- `load_ckpt(path)`：从 checkpoint 中提取 `splats`（means/scales/quats/...）。
+- `load_ply_sequence(path | list)`：加载单帧或序列 PLY，返回统一的张量字典。
+- 两者都转换成标准的 `Dict[str, Tensor]` 供后续模块使用。
+
+### 3. 编码流水线模块
+- `PruningStage`：离群点过滤、mask 过滤等，可配置开关。
+- `QuantizationStage`：基于属性设置 bitwidth、clamp range，复用现有 `_compress_*` 逻辑。
+- `MappingStage`：排序/映射策略（PLAS、morton、无排序）；负责 3D→2D 重排。
+- `CodecStage`：调度 `gsplat/compression` 中的具体 codec 类。
+- 元信息统一写入 `meta.json`（沿用现状）。
+
+### 4. 解码流水线模块
+- `CodecStage.decode`：调用 codec 的 `decompress`，获得属性张量。
+- `DequantizationStage`：执行逆变换、逆量化；可复用 `inverse_log_transform` 等函数。
+- 最终输出：
+  - 内存中的 `Dict[str, Tensor]`，或
+  - 调用 `save_ply` 写入磁盘。
+
+### 5. 配置结构
+使用 dataclass 组织配置，便于 YAML/CLI：
+```python
+@dataclass
+class PTCompressionConfig:
+    input_type: Literal["ply", "ckpt"]
+    codec: Literal["png", "entropy", "hevc", "seq_hevc", "seq_yuv"]
+    codec_params: Dict[str, Any] = field(default_factory=dict)
+    preprocess: PreprocessConfig = field(default_factory=PreprocessConfig)
+    quant: QuantConfig = field(default_factory=QuantConfig)
+```
+其中 `PreprocessConfig`、`QuantConfig` 再细分字段（是否过滤、排序策略、bitwidth 等）。
+
+### 6. CLI/脚本整合
+- 重写/新增脚本（如 `examples/benchmarks/post_train_compress.sh`）调用 `PostTrainingCompressor`。
+- YAML/CLI 用同一套配置，避免脚本内硬编码。
+- 脚本流程：解析配置 → `encode()` → `decode()` → 统计评估 → 输出路径。
+
+## 顾及到的补充需求
+- 模块内部独立提供 PLY / CKPT 的读写接口（与 encode/decode 解耦）。
+- 支持读取单帧或序列 splats，保留 `is_sequence` 或 `num_frames` 标识。
+- `encode()` 前和 `decode()` 后输出参数分布直方图（matplotlib），但不写入 meta 数据。
+- 实现时参考现有 `gsplat/compression` 的操作顺序，避免与现有流程背离。
+
+## `PostTrainingComp` 草图
+```python
+@dataclass
+class PostTrainingComp:
+    input_spec: InputSpec
+    codec_config: CodecConfig
+    preprocess_cfg: PreprocessConfig
+    quant_cfg: QuantConfig
+    output_dir: Path
+
+    splats: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = field(init=False)
+    is_sequence: bool = field(init=False)
+    codec: BaseCodec = field(init=False)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+    def load_inputs(self) -> None:
+        if self.input_spec.type == "ply":
+            self.splats = load_ply(self.input_spec.path, as_sequence=self.input_spec.as_sequence)
+        elif self.input_spec.type == "ckpt":
+            self.splats = load_ckpt(self.input_spec.path)
+        self.is_sequence = isinstance(self.splats, list)
+        self.metadata["num_frames"] = len(self.splats) if self.is_sequence else 1
+
+    def encode(self) -> CompressionResult:
+        self.load_inputs()
+        self._plot_stats(self.splats, stage="before_encode")
+        payload = self._run_encode_pipeline(self.splats)
+        self._save_payload(payload)
+        return CompressionResult(payload_path=..., metadata=self.metadata)
+
+    def decode(self, payload_path: Path) -> DecodeResult:
+        payload = self._load_payload(payload_path)
+        decoded = self._run_decode_pipeline(payload)
+        self._plot_stats(decoded, stage="after_decode")
+        self._write_outputs(decoded)
+        return DecodeResult(splats=decoded, saved_paths=...)
+
+    def _run_encode_pipeline(self, splats):
+        pruned = run_pruning(splats, self.preprocess_cfg)
+        quantized = run_quantization(pruned, self.quant_cfg)
+        mapped, mapping_ctx = run_mapping(quantized, self.preprocess_cfg.mapping)
+        encoded = self.codec.encode(mapped, context=mapping_ctx)
+        return {"encoded": encoded, "mapping_ctx": mapping_ctx}
+
+    def _run_decode_pipeline(self, payload):
+        decoded = self.codec.decode(payload["encoded"], context=payload.get("mapping_ctx"))
+        unmapped = run_inverse_mapping(decoded, payload.get("mapping_ctx"))
+        dequant = run_dequantization(unmapped, self.quant_cfg)
+        restored = run_postprocess(dequant, self.preprocess_cfg)
+        return restored
+
+    def _plot_stats(self, splats, stage: str) -> None:
+        # 遍历属性画直方图，保存到 output_dir/stage_* 下
+```
+
+- `InputSpec`、`CodecConfig`、`PreprocessConfig`、`QuantConfig` 等 dataclass 可进一步细化。
+- `CompressionResult` / `DecodeResult` 用于统一返回路径、元数据。
+- PLY/CKPT 的读写实现为独立的 util 函数。
+- `BaseCodec` 为各 codec 的抽象基类，具体实现参考现有 PNG/Entropy/HEVC 等类。
+
+## 未决问题
+1. 是否需要一次性处理多个场景/帧并输出统一码流？
+2. 统计文件命名、保存位置是否需要可配置？
+3. 与现有脚本结合时，是否要提供默认 YAML 模板。
+
+## 设计补充（2024-xx-xx）
+- `Quantization` / `Dequantization` 阶段为必选流程，需始终执行。
+- `Pruning`、`Mapping`、`Codec` 阶段可以按需求启用或跳过，可通过配置显式控制。
+- `PostTrainingComp` 需提供 `save_decoded_splats(decoded, destination)` 助手，用于在 `decode()` 完成后将还原的 splats 持久化到 PLY / CKPT 等目标格式，避免重复实现写盘逻辑。

examples/benchmarks/static_gscodec/mcmc_tt_sim.sh

@@ -0,0 +1,97 @@
+# ----------------- Training Setting-------------- #
+SCENE_DIR="data/tandt"
+# eval all 9 scenes for benchmarking
+SCENE_LIST="train truck" # train truck
+# SCENE_LIST="garden bicycle stump bonsai counter kitchen room treehill flowers"
+
+# # 0.36M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_0_36M_png_compression"
+# CAP_MAX=360000
+
+# # 0.49M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_tt_0_49M_png_compression"
+# CAP_MAX=490000
+
+# 1M GSs
+RESULT_DIR="results/new_cfg_tt"
+CAP_MAX=1000000
+
+# # 4M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_4M_png_compression"
+# CAP_MAX=4000000
+
+RD_LAMBDA=0.01
+
+# ----------------- Training Setting-------------- #
+
+# ----------------- Args ------------------------- #
+
+if [ ! -z "$1" ]; then
+    RD_LAMBDA="$1"
+    RESULT_DIR="results/new_cfg_tt_rd_lambda_${RD_LAMBDA}"
+fi
+
+# ----------------- Args ------------------------- #
+
+# ----------------- Main Job --------------------- #
+run_single_scene() {
+    local GPU_ID=$1
+    local SCENE=$2
+
+    echo "Running $SCENE on GPU: $GPU_ID"
+
+    # train without eval
+    CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --eval_steps -1 --disable_viewer --data_factor 1 \
+        --strategy.cap-max $CAP_MAX \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/ \
+        --compression_sim \
+        --entropy_model_opt \
+        --rd_lambda $RD_LAMBDA \
+        --shN_ada_mask_opt \
+        --compression png
+
+
+    # eval: use vgg for lpips to align with other benchmarks
+    CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --disable_viewer --data_factor 1 \
+        --strategy.cap-max $CAP_MAX \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/ \
+        --lpips_net vgg \
+        --ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt \
+        --compression png
+    
+}
+# ----------------- Main Job --------------------- #
+
+
+
+# ----------------- Experiment Loop -------------- #
+GPU_LIST=(6 7)
+GPU_COUNT=${#GPU_LIST[@]}
+
+SCENE_IDX=-1
+
+for SCENE in $SCENE_LIST;
+do
+    SCENE_IDX=$((SCENE_IDX + 1))
+    {
+        run_single_scene ${GPU_LIST[$SCENE_IDX]} $SCENE
+    } #&
+
+done
+
+# ----------------- Experiment Loop -------------- #
+
+# Wait for finishing the jobs across all scenes 
+wait
+echo "All scenes finished."
+
+# Zip the compressed files and summarize the stats
+if command -v zip &> /dev/null
+then
+    echo "Zipping results"
+    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
+else
+    echo "zip command not found, skipping zipping"
+fi