Star53
9
代码
代码
Issues9
Pull Requests3
Actions
项目讨论
Wiki
分析
Star53
9
已合并
feat: 3DGS densification (DefaultStrategy) #14
xieyajun创建于 29 天前
feat: 3DGS densification (DefaultStrategy) #14
已合并
xieyajun创建于 29 天前
从xieyajun/gauss-splatfeat/densify合入到cann/gauss-splatmaster
共 12 个文件变更+869-112
Aexamples/gsplat/__init__.py+10-0
@@ -0,0 +1,10 @@
1+# coding: utf-8
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+ 
10+ 
Aexamples/gsplat/strategy/__init__.py+11-0
@@ -0,0 +1,11 @@
1+# coding: utf-8
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+ 
10+from .base import Strategy
11+from .default import DefaultStrategy
Aexamples/gsplat/strategy/base.py+59-0
@@ -0,0 +1,59 @@
1+# coding: utf-8
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+ 
10+from dataclasses import dataclass
11+from typing import Dict, Union
12+ 
13+import torch
14+ 
15+ 
16+@dataclass
17+class Strategy:
18+ """Base class for the GS densification strategy.
19+ 
20+ This class is an base class that defines the interface for the GS
21+ densification strategy.
22+ """
23+ 
24+ def check_sanity(
25+ self,
26+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
27+ optimizers: Dict[str, torch.optim.Optimizer],
28+ ):
29+ """Sanity check for the parameters and optimizers."""
30+ trainable_params = set(
31+ [name for name, param in params.items() if param.requires_grad]
32+ )
33+ assert trainable_params == set(optimizers.keys()), (
34+ "trainable parameters and optimizers must have the same keys, "
35+ f"but got {trainable_params} and {optimizers.keys()}"
36+ )
37+ 
38+ for optimizer in optimizers.values():
39+ assert len(optimizer.param_groups) == 1, (
40+ "Each optimizer must have exactly one param_group, "
41+ "that cooresponds to each parameter, "
42+ f"but got {len(optimizer.param_groups)}"
43+ )
44+ 
45+ def step_pre_backward(
46+ self,
47+ *args,
48+ **kwargs,
49+ ):
50+ """Callback function to be executed before the `loss.backward()` call."""
51+ pass
52+ 
53+ def step_post_backward(
54+ self,
55+ *args,
56+ **kwargs,
57+ ):
58+ """Callback function to be executed after the `loss.backward()` call."""
59+ pass
Aexamples/gsplat/strategy/default.py+385-0
@@ -0,0 +1,385 @@
1+# coding: utf-8
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+ 
10+from dataclasses import dataclass
11+from typing import Any, Dict, Tuple, Union
12+ 
13+import torch
14+from typing_extensions import Literal
15+ 
16+from .base import Strategy
17+from .ops import duplicate, remove, reset_opa, split
18+ 
19+ 
20+@dataclass
21+class DefaultStrategy(Strategy):
22+ """A default strategy that follows the original 3DGS paper:
23+ 
24+ `3D Gaussian Splatting for Real-Time Radiance Field Rendering <https://arxiv.org/abs/2308.04079>`_
25+ 
26+ The strategy will:
27+ 
28+ - Periodically duplicate GSs with high image plane gradients and small scales.
29+ - Periodically split GSs with high image plane gradients and large scales.
30+ - Periodically prune GSs with low opacity.
31+ - Periodically reset GSs to a lower opacity.
32+ 
33+ If `absgrad=True`, it will use the absolute gradients instead of average gradients
34+ for GS duplicating & splitting, following the AbsGS paper:
35+ 
36+ `AbsGS: Recovering Fine Details for 3D Gaussian Splatting <https://arxiv.org/abs/2404.10484>`_
37+ 
38+ Which typically leads to better results but requires to set the `grow_grad2d` to a
39+ higher value, e.g., 0.0008. Also, the :func:`rasterization` function should be called
40+ with `absgrad=True` as well so that the absolute gradients are computed.
41+ 
42+ Args:
43+ prune_opa (float): GSs with opacity below this value will be pruned. Default is 0.005.
44+ grow_grad2d (float): GSs with image plane gradient above this value will be
45+ split/duplicated. Default is 0.0002.
46+ grow_scale3d (float): GSs with 3d scale (normalized by scene_scale) below this
47+ value will be duplicated. Above will be split. Default is 0.01.
48+ grow_scale2d (float): GSs with 2d scale (normalized by image resolution) above
49+ this value will be split. Default is 0.05.
50+ prune_scale3d (float): GSs with 3d scale (normalized by scene_scale) above this
51+ value will be pruned. Default is 0.1.
52+ prune_scale2d (float): GSs with 2d scale (normalized by image resolution) above
53+ this value will be pruned. Default is 0.15.
54+ refine_scale2d_stop_iter (int): Stop refining GSs based on 2d scale after this
55+ iteration. Default is 0. Set to a positive value to enable this feature.
56+ refine_start_iter (int): Start refining GSs after this iteration. Default is 500.
57+ refine_stop_iter (int): Stop refining GSs after this iteration. Default is 15_000.
58+ reset_every (int): Reset opacities every this steps. Default is 3000.
59+ refine_every (int): Refine GSs every this steps. Default is 100.
60+ pause_refine_after_reset (int): Pause refining GSs until this number of steps after
61+ reset, Default is 0 (no pause at all) and one might want to set this number to the
62+ number of images in training set.
63+ absgrad (bool): Use absolute gradients for GS splitting. Default is False.
64+ revised_opacity (bool): Whether to use revised opacity heuristic from
65+ arXiv:2404.06109 (experimental). Default is False.
66+ verbose (bool): Whether to print verbose information. Default is False.
67+ key_for_gradient (str): Which variable uses for densification strategy.
68+ 3DGS uses "means2d" gradient and 2DGS uses a similar gradient which stores
69+ in variable "gradient_2dgs".
70+ 
71+ Examples:
72+ 
73+ >>> from gsplat import DefaultStrategy, rasterization
74+ >>> params: Dict[str, torch.nn.Parameter] | torch.nn.ParameterDict = ...
75+ >>> optimizers: Dict[str, torch.optim.Optimizer] = ...
76+ >>> strategy = DefaultStrategy()
77+ >>> strategy.check_sanity(params, optimizers)
78+ >>> strategy_state = strategy.initialize_state()
79+ >>> for step in range(1000):
80+ ... render_image, render_alpha, info = rasterization(...)
81+ ... strategy.step_pre_backward(params, optimizers, strategy_state, step, info)
82+ ... loss = ...
83+ ... loss.backward()
84+ ... strategy.step_post_backward(params, optimizers, strategy_state, step, info)
85+ 
86+ """
87+ 
88+ prune_opa: float = 0.005
89+ grow_grad2d: float = 0.0002
90+ grow_scale3d: float = 0.01
91+ grow_scale2d: float = 0.05
92+ prune_scale3d: float = 0.1
93+ prune_scale2d: float = 0.15
94+ refine_scale2d_stop_iter: int = 0
95+ refine_start_iter: int = 500
96+ refine_stop_iter: int = 15_000
97+ reset_every: int = 3000
98+ refine_every: int = 100
99+ pause_refine_after_reset: int = 0
100+ # hard cap on number of GSs (0 = off). NEEDED for outdoor scenes: once gradient
101+ # attribution is correct (proj_filter fix), growth is aggressive and can overshoot
102+ # the AscendC kernel's healthy range (~3.7M GS -> ERR99999) -> quality collapse.
103+ cap_max: int = 3_000_000
104+ # screen-space pruning: track per-GS max screen radius and prune artifacts (big-on-screen
105+ # GSs). On NPU this also avoids degenerate gaussians that trigger kernel ERR99999 on
106+ # outdoor scenes. Decoupled from screen-space SPLIT (which stays off) to avoid overgrowth.
107+ screen_prune: bool = True
108+ # opacity reset = UPPER CAP (anti-saturation, prevents over-confident GSs/floaters).
109+ # NOT gsplat's force-low-to-0.01 (which blacks out the render and the NPU backward
110+ # can't recover -> collapse). High-ish cap (0.1) is gentle + needed for outdoor quality.
111+ reset_value: float = 0.1
112+ absgrad: bool = False
113+ revised_opacity: bool = False
114+ verbose: bool = False
115+ key_for_gradient: Literal["means2d", "gradient_2dgs"] = "means2d"
116+ 
117+ def initialize_state(self, scene_scale: float = 1.0) -> Dict[str, Any]:
118+ """Initialize and return the running state for this strategy.
119+ 
120+ The returned state should be passed to the `step_pre_backward()` and
121+ `step_post_backward()` functions.
122+ """
123+ # Postpone the initialization of the state to the first step so that we can
124+ # put them on the correct device.
125+ # - grad2d: running accum of the norm of the image plane gradients for each GS.
126+ # - count: running accum of how many time each GS is visible.
127+ # - radii: the radii of the GSs (normalized by the image resolution).
128+ state = {"grad2d": None, "count": None, "radii": None, "scene_scale": scene_scale}
129+ return state
130+ 
131+ def check_sanity(
132+ self,
133+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
134+ optimizers: Dict[str, torch.optim.Optimizer],
135+ ):
136+ """Sanity check for the parameters and optimizers.
137+ 
138+ Check if:
139+ * `params` and `optimizers` have the same keys.
140+ * Each optimizer has exactly one param_group, corresponding to each parameter.
141+ * The following keys are present: {"means", "scales", "quats", "opacities"}.
142+ 
143+ Raises:
144+ AssertionError: If any of the above conditions is not met.
145+ 
146+ .. note::
147+ It is not required but highly recommended for the user to call this function
148+ after initializing the strategy to ensure the convention of the parameters
149+ and optimizers is as expected.
150+ """
151+ 
152+ super().check_sanity(params, optimizers)
153+ # The following keys are required for this strategy.
154+ for key in ["means", "scales", "quats", "opacities"]:
155+ assert key in params, f"{key} is required in params but missing."
156+ 
157+ def step_pre_backward(
158+ self,
159+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
160+ optimizers: Dict[str, torch.optim.Optimizer],
161+ state: Dict[str, Any],
162+ step: int,
163+ info: Dict[str, Any],
164+ ):
165+ """Callback function to be executed before the `loss.backward()` call."""
166+ assert (
167+ self.key_for_gradient in info
168+ ), "The 2D means of the Gaussians is required but missing."
169+ info[self.key_for_gradient].retain_grad()
170+ 
171+ def step_post_backward(
172+ self,
173+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
174+ optimizers: Dict[str, torch.optim.Optimizer],
175+ state: Dict[str, Any],
176+ step: int,
177+ info: Dict[str, Any],
178+ packed: bool = False,
179+ ):
180+ """Callback function to be executed after the `loss.backward()` call."""
181+ if step >= self.refine_stop_iter:
182+ return
183+ 
184+ self._update_state(params, state, info, packed=packed)
185+ 
186+ if (
187+ step > self.refine_start_iter
188+ and step % self.refine_every == 0
189+ and step % self.reset_every >= self.pause_refine_after_reset
190+ ):
191+ # grow GSs
192+ n_dupli, n_split = self._grow_gs(params, optimizers, state, step)
193+ if self.verbose:
194+ print(
195+ f"Step {step}: {n_dupli} GSs duplicated, {n_split} GSs split. "
196+ f"Now having {len(params['means'])} GSs."
197+ )
198+ 
199+ # prune GSs
200+ n_prune = self._prune_gs(params, optimizers, state, step)
201+ if self.verbose:
202+ print(
203+ f"Step {step}: {n_prune} GSs pruned. "
204+ f"Now having {len(params['means'])} GSs."
205+ )
206+ 
207+ # hard cap: if over cap_max, remove the lowest-opacity excess GSs
208+ if self.cap_max > 0 and len(params["means"]) > self.cap_max:
209+ n_excess = len(params["means"]) - self.cap_max
210+ opa = torch.sigmoid(params["opacities"].flatten())
211+ _, low_ids = torch.topk(opa, k=n_excess, largest=False)
212+ cap_mask = torch.zeros(len(opa), dtype=torch.bool, device=opa.device)
213+ cap_mask[low_ids] = True
214+ remove(params=params, optimizers=optimizers, state=state, mask=cap_mask)
215+ if self.verbose:
216+ print(f"Step {step}: capped to {self.cap_max} GSs (removed {n_excess}).")
217+ 
218+ # reset running stats
219+ state["grad2d"].zero_()
220+ state["count"].zero_()
221+ state["radii"].zero_()
222+ torch.cuda.empty_cache()
223+ 
224+ # opacity reset as anti-saturation CAP (not force-low). Fixed `&`->`and` precedence
225+ # bug so it actually fires; value=reset_value caps opacity to prevent saturation.
226+ if step % self.reset_every == 0 and step > 0:
atomgit-bot
atomgit-botatomgit-bot29 天前

🟡 Medium Priority

changed line: if step >= self.refine_stop_iter: return（第181-182行）→ 受影响行为：step_post_backward 中的提前返回同时跳过了 _update_state 和 reset_opa（第226行）→ 失败模式：在 refine_stop_iter（默认15,000）之后，不透明度重置完全停止。在标准 3DGS 中，不透明度重置在 refine 停止后会持续整个训练过程（例如直到第30,000步），以防止漂浮物/过度饱和。在此实现中，对于 max_steps=30,000，训练的整个后半段都没有不透明度上限重置，可能导致质量下降（过度自信的高斯、漂浮物）。注意：runner.py 在每一步都无条件调用 step_post_backward，因此在 refine_stop_iter 之后该提前返回必然触发。

建议：将不透明度重置逻辑移到提前返回之前，或在提前返回中复制一份重置调用，使不透明度上限在 refine_stop_iter 之后继续生效。

改动建议
226
+ if step >= self.refine_stop_iter:
227
+ # still apply periodic opacity reset even after refinement stops
226
- if step % self.reset_every == 0 and step > 0:
228
+ if step % self.reset_every == 0 and step > 0:
229
+ reset_opa(
230
+ params=params,
231
+ optimizers=optimizers,
232
+ state=state,
233
+ value=self.reset_value,
234
+ )
235
+ return
应用建议
likedislike
227+ reset_opa(
228+ params=params,
229+ optimizers=optimizers,
230+ state=state,
231+ value=self.reset_value,
232+ )
233+ 
234+ def _update_state(
235+ self,
236+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
237+ state: Dict[str, Any],
238+ info: Dict[str, Any],
239+ packed: bool = False,
240+ ):
241+ for key in [
242+ "width",
243+ "height",
244+ "n_cameras",
245+ "radii",
246+ "gaussian_ids",
247+ self.key_for_gradient,
248+ ]:
249+ assert key in info, f"{key} is required but missing."
250+ 
251+ # normalize grads to [-1, 1] screen space
252+ if self.absgrad:
atomgit-bot
atomgit-botatomgit-bot29 天前

🟡 Medium Priority

changed line: grads = info[self.key_for_gradient].absgrad.clone() (第253行) → 行为：当 absgrad=True 时访问 .absgrad 属性 → 失败模式：代码库中没有任何地方在 tensor 上设置 .absgrad 属性。rendering.py 的 rasterization 流程不会创建该属性，runner.py 也只使用 .grad。因此，任何将 absgrad=True 传入 DefaultStrategy 的尝试都会立即引发 AttributeError: 'Tensor' object has no attribute 'absgrad'。当前默认值为 False，因此未被触发，但这是潜伏的运行时可触发缺陷。

建议：要么移除 absgrad 支持（如果 NPU 渲染管线不支持），要么在 rasterization 输出中正确填充 .absgrad 属性。临时缓解方案：回退到 .grad 或在缺少时抛出明确错误。

改动建议
252
- if self.absgrad:
252
+ if self.absgrad:
253
+ # absgrad attribute not currently set by NPU rasterization;
254
+ # fall back to .grad (user must also compute abs grad externally).
255
+ grads = info[self.key_for_gradient].grad.abs().clone() if hasattr(info[self.key_for_gradient], 'absgrad') else info[self.key_for_gradient].grad.clone()
256
+ else:
应用建议
likedislike
253+ grads = info[self.key_for_gradient].absgrad.clone()
254+ else:
255+ grads = info[self.key_for_gradient].grad.clone()
256+ grads[..., 0] *= info["width"] / 2.0 * info["n_cameras"]
257+ grads[..., 1] *= info["height"] / 2.0 * info["n_cameras"]
258+ 
259+ # initialize state on the first run
260+ n_gaussian = len(list(params.values())[0])
261+ 
262+ if state["grad2d"] is None:
263+ state["grad2d"] = torch.zeros(n_gaussian, device=grads.device)
264+ if state["count"] is None:
265+ state["count"] = torch.zeros(n_gaussian, device=grads.device)
266+ if state["radii"] is None:
267+ state["radii"] = torch.zeros(n_gaussian, device=grads.device)
268+ 
269+ # update the running state
270+ if packed:
271+ # grads is [nnz, 2]
272+ gs_ids = info["gaussian_ids"] # [nnz]
273+ radii = info["radii"].max(dim=-1).values # [nnz]
274+ else:
275+ # [proj_filter fix — the ONE wiring change] means2d/radii/grad are COMPACTED
276+ # (visible GSs packed in front, original-index ascending order). where(radii>0)
277+ # gives compacted positions, NOT original gaussian ids -> grads misattributed.
278+ # Decode proj_filter (packed bitfield over original N) to recover true ids.
279+ pf = info.get("proj_filter", None)
280+ if pf is not None:
281+ import numpy as np
282+ n_gs = int(info.get("n_gaussians", state["grad2d"].numel()))
283+ bits = np.unpackbits(pf.reshape(-1).cpu().numpy())[:n_gs].astype(bool)
284+ vis = torch.from_numpy(bits).to(grads.device)
285+ gs_ids = vis.nonzero(as_tuple=True)[0]
286+ sel = (info["radii"] > 0.0).all(dim=-1)
287+ grads_sel = grads[sel]
288+ radii_sel = info["radii"][sel].max(dim=-1).values
289+ nvis = min(gs_ids.numel(), grads_sel.shape[0])
290+ gs_ids = gs_ids[:nvis]
291+ grads = grads_sel[:nvis]
292+ radii = radii_sel[:nvis]
293+ else:
294+ sel = (info["radii"] > 0.0).all(dim=-1) # [C, N]
295+ gs_ids = torch.where(sel)[1] # [nnz]
296+ grads = grads[sel] # [nnz, 2]
297+ radii = info["radii"][sel].max(dim=-1).values # [nnz]
298+ state["grad2d"].index_add_(0, gs_ids, grads.norm(dim=-1))
299+ state["count"].index_add_(
300+ 0, gs_ids, torch.ones_like(gs_ids, dtype=torch.float32)
301+ )
302+ # always track per-GS max screen radius (normalized) for screen-space pruning
303+ state["radii"][gs_ids] = torch.maximum(
304+ state["radii"][gs_ids],
305+ radii / float(max(info["width"], info["height"])),
306+ )
307+ 
308+ @torch.no_grad()
309+ def _grow_gs(
310+ self,
311+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
312+ optimizers: Dict[str, torch.optim.Optimizer],
313+ state: Dict[str, Any],
314+ step: int,
315+ ) -> Tuple[int, int]:
316+ count = state["count"]
317+ grads = state["grad2d"] / count.clamp_min(1)
318+ device = grads.device
319+ 
320+ is_grad_high = grads > self.grow_grad2d
321+ is_small = (
322+ torch.exp(params["scales"]).max(dim=-1).values
323+ <= self.grow_scale3d * state["scene_scale"]
324+ )
325+ is_dupli = is_grad_high & is_small
326+ n_dupli = is_dupli.sum().item()
327+ 
328+ is_large = ~is_small
329+ is_split = is_grad_high & is_large
330+ if step < self.refine_scale2d_stop_iter:
331+ is_split |= state["radii"] > self.grow_scale2d
332+ n_split = is_split.sum().item()
333+ 
334+ # first duplicate
335+ if n_dupli > 0:
336+ duplicate(params=params, optimizers=optimizers, state=state, mask=is_dupli)
337+ 
338+ # new GSs added by duplication will not be split
339+ is_split = torch.cat(
340+ [
341+ is_split,
342+ torch.zeros(n_dupli, dtype=torch.bool, device=device),
343+ ]
344+ )
345+ 
346+ # then split
347+ if n_split > 0:
348+ split(
349+ params=params,
350+ optimizers=optimizers,
351+ state=state,
352+ mask=is_split,
353+ revised_opacity=self.revised_opacity,
354+ )
355+ return n_dupli, n_split
356+ 
357+ @torch.no_grad()
358+ def _prune_gs(
359+ self,
360+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
361+ optimizers: Dict[str, torch.optim.Optimizer],
362+ state: Dict[str, Any],
363+ step: int,
364+ ) -> int:
365+ is_prune = torch.sigmoid(params["opacities"].flatten()) < self.prune_opa
366+ # screen-space pruning (decoupled from screen-split, ungated): prune artifacts /
367+ # degenerate big-on-screen GSs. On NPU this also avoids the kernel ERR99999 that
368+ # otherwise hits on outdoor scenes once growth is correct (proj_filter fix).
369+ if self.screen_prune and state.get("radii") is not None:
370+ is_prune = is_prune | (state["radii"] > self.prune_scale2d)
371+ if step > self.reset_every:
372+ is_too_big = (
373+ torch.exp(params["scales"]).max(dim=-1).values
374+ > self.prune_scale3d * state["scene_scale"]
375+ )
376+ if step < self.refine_scale2d_stop_iter:
377+ is_too_big |= state["radii"] > self.prune_scale2d
378+ 
379+ is_prune = is_prune | is_too_big
380+ 
381+ n_prune = is_prune.sum().item()
382+ if n_prune > 0:
383+ remove(params=params, optimizers=optimizers, state=state, mask=is_prune)
384+ 
385+ return n_prune
Aexamples/gsplat/strategy/ops.py+213-0
@@ -0,0 +1,213 @@
1+# coding: utf-8
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+ 
10+from typing import Callable, Dict, List, Union
11+ 
12+import torch
13+import torch.nn.functional as F
14+from torch import Tensor
15+ 
16+from gsplat.utils import normalized_quat_to_rotmat
17+ 
18+ 
19+@torch.no_grad()
20+def _update_param_with_optimizer(
21+ param_fn: Callable[[str, Tensor], Tensor],
22+ optimizer_fn: Callable[[str, Tensor], Tensor],
23+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
24+ optimizers: Dict[str, torch.optim.Optimizer],
25+ names: Union[List[str], None] = None,
26+):
27+ """Update the parameters and the state in the optimizers with defined functions.
28+ 
29+ Args:
30+ param_fn: A function that takes the name of the parameter and the parameter itself,
31+ and returns the new parameter.
32+ optimizer_fn: A function that takes the key of the optimizer state and the state value,
33+ and returns the new state value.
34+ params: A dictionary of parameters.
35+ optimizers: A dictionary of optimizers, each corresponding to a parameter.
36+ names: A list of key names to update. If None, update all. Default: None.
37+ """
38+ if names is None:
39+ # If names is not provided, update all parameters
40+ names = list(params.keys())
41+ 
42+ for name in names:
43+ param = params[name]
44+ new_param = param_fn(name, param)
45+ params[name] = new_param
46+ if name not in optimizers:
47+ assert not param.requires_grad, (
48+ f"Optimizer for {name} is not found, but the parameter is trainable."
49+ f"Got requires_grad={param.requires_grad}"
50+ )
51+ continue
52+ optimizer = optimizers[name]
53+ for i in range(len(optimizer.param_groups)):
54+ param_state = optimizer.state[param]
55+ del optimizer.state[param]
56+ for key in param_state.keys():
57+ if key != "step":
58+ v = param_state[key]
59+ param_state[key] = optimizer_fn(key, v)
60+ optimizer.param_groups[i]["params"] = [new_param]
61+ optimizer.state[new_param] = param_state
62+ 
63+ 
64+@torch.no_grad()
65+def duplicate(
66+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
67+ optimizers: Dict[str, torch.optim.Optimizer],
68+ state: Dict[str, Tensor],
69+ mask: Tensor,
70+):
71+ """Inplace duplicate the Gaussian with the given mask.
72+ 
73+ Args:
74+ params: A dictionary of parameters.
75+ optimizers: A dictionary of optimizers, each corresponding to a parameter.
76+ mask: A boolean mask to duplicate the Gaussians.
77+ """
78+ device = mask.device
79+ sel = torch.where(mask)[0]
80+ 
81+ def param_fn(name: str, p: Tensor) -> Tensor:
82+ return torch.nn.Parameter(torch.cat([p, p[sel]]), requires_grad=p.requires_grad)
83+ 
84+ def optimizer_fn(key: str, v: Tensor) -> Tensor:
85+ return torch.cat([v, torch.zeros((len(sel), *v.shape[1:]), device=device)])
86+ 
87+ # update the parameters and the state in the optimizers
88+ _update_param_with_optimizer(param_fn, optimizer_fn, params, optimizers)
89+ # update the extra running state
90+ for k, v in state.items():
91+ if isinstance(v, torch.Tensor):
92+ state[k] = torch.cat((v, v[sel]))
93+ 
94+ 
95+@torch.no_grad()
96+def split(
97+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
98+ optimizers: Dict[str, torch.optim.Optimizer],
99+ state: Dict[str, Tensor],
100+ mask: Tensor,
101+ revised_opacity: bool = False,
102+):
103+ """Inplace split the Gaussian with the given mask.
104+ 
105+ Args:
106+ params: A dictionary of parameters.
107+ optimizers: A dictionary of optimizers, each corresponding to a parameter.
108+ mask: A boolean mask to split the Gaussians.
109+ revised_opacity: Whether to use revised opacity formulation
110+ from arXiv:2404.06109. Default: False.
111+ """
112+ device = mask.device
113+ sel = torch.where(mask)[0]
114+ rest = torch.where(~mask)[0]
115+ 
116+ scales = torch.exp(params["scales"][sel])
117+ quats = F.normalize(params["quats"][sel], dim=-1)
118+ rotmats = normalized_quat_to_rotmat(quats) # [N, 3, 3]
119+ samples = torch.einsum(
120+ "nij,nj,bnj->bni",
121+ rotmats,
122+ scales,
123+ torch.randn(2, len(scales), 3, device=device),
124+ ) # [2, N, 3]
125+ 
126+ def param_fn(name: str, p: Tensor) -> Tensor:
127+ repeats = [2] + [1] * (p.dim() - 1)
128+ if name == "means":
129+ p_split = (p[sel] + samples).reshape(-1, 3) # [2N, 3]
130+ elif name == "scales":
131+ p_split = torch.log(scales / 1.6).repeat(2, 1) # [2N, 3]
132+ elif name == "opacities" and revised_opacity:
133+ new_opacities = 1.0 - torch.sqrt(1.0 - torch.sigmoid(p[sel]))
134+ p_split = torch.logit(new_opacities).repeat(repeats) # [2N]
135+ else:
136+ p_split = p[sel].repeat(repeats)
137+ p_new = torch.cat([p[rest], p_split])
138+ p_new = torch.nn.Parameter(p_new, requires_grad=p.requires_grad)
139+ return p_new
140+ 
141+ def optimizer_fn(key: str, v: Tensor) -> Tensor:
142+ v_split = torch.zeros((2 * len(sel), *v.shape[1:]), device=device)
143+ return torch.cat([v[rest], v_split])
144+ 
145+ # update the parameters and the state in the optimizers
146+ _update_param_with_optimizer(param_fn, optimizer_fn, params, optimizers)
147+ # update the extra running state
148+ for k, v in state.items():
149+ if isinstance(v, torch.Tensor):
150+ repeats = [2] + [1] * (v.dim() - 1)
151+ v_new = v[sel].repeat(repeats)
152+ state[k] = torch.cat((v[rest], v_new))
153+ 
154+ 
155+@torch.no_grad()
156+def remove(
157+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
158+ optimizers: Dict[str, torch.optim.Optimizer],
159+ state: Dict[str, Tensor],
160+ mask: Tensor,
161+):
162+ """Inplace remove the Gaussian with the given mask.
163+ 
164+ Args:
165+ params: A dictionary of parameters.
166+ optimizers: A dictionary of optimizers, each corresponding to a parameter.
167+ mask: A boolean mask to remove the Gaussians.
168+ """
169+ sel = torch.where(~mask)[0]
170+ 
171+ def param_fn(name: str, p: Tensor) -> Tensor:
172+ return torch.nn.Parameter(p[sel], requires_grad=p.requires_grad)
173+ 
174+ def optimizer_fn(key: str, v: Tensor) -> Tensor:
175+ return v[sel]
176+ 
177+ # update the parameters and the state in the optimizers
178+ _update_param_with_optimizer(param_fn, optimizer_fn, params, optimizers)
179+ # update the extra running state
180+ for k, v in state.items():
181+ if isinstance(v, torch.Tensor):
182+ state[k] = v[sel]
183+ 
184+ 
185+@torch.no_grad()
186+def reset_opa(
187+ params: Union[Dict[str, torch.nn.Parameter], torch.nn.ParameterDict],
188+ optimizers: Dict[str, torch.optim.Optimizer],
189+ state: Dict[str, Tensor],
190+ value: float,
191+):
192+ """Inplace reset the opacities to the given post-sigmoid value.
193+ 
194+ Args:
195+ params: A dictionary of parameters.
196+ optimizers: A dictionary of optimizers, each corresponding to a parameter.
197+ value: The value to reset the opacities
198+ """
199+ 
200+ def param_fn(name: str, p: Tensor) -> Tensor:
201+ if name == "opacities":
202+ opacities = torch.clamp(p, max=torch.logit(torch.tensor(value, device=p.device)).item())
203+ return torch.nn.Parameter(opacities, requires_grad=p.requires_grad)
204+ else:
205+ raise ValueError(f"Unexpected parameter name: {name}")
206+ 
207+ def optimizer_fn(key: str, v: Tensor) -> Tensor:
208+ return torch.zeros_like(v)
209+ 
210+ # update the parameters and the state in the optimizers
211+ _update_param_with_optimizer(
212+ param_fn, optimizer_fn, params, optimizers, names=["opacities"]
213+ )
Mexamples/gsplat/utils.py+12-1
@@ -107,4 +107,15 @@ def save_ply(splats: torch.nn.ParameterDict, dir: str, colors: torch.Tensor = No
107 107 
108 for data in [scales, quats]:108 for data in [scales, quats]:
109 for j in range(data.shape[1]):109 for j in range(data.shape[1]):
110- f.write(struct.pack("<f", data[i, j]))110+ f.write(struct.pack("<f", data[i, j]))
111+ 
112+ 
113+def normalized_quat_to_rotmat(quat):
114+ """Quaternion (w,x,y,z, normalized) -> 3x3 rotation matrix. Pure-python NPU-friendly."""
115+ w, x, y, z = quat.unbind(-1)
116+ R = torch.stack([
117+ 1 - 2*(y*y + z*z), 2*(x*y - w*z), 2*(x*z + w*y),
118+ 2*(x*y + w*z), 1 - 2*(x*x + z*z), 2*(y*z - w*x),
119+ 2*(x*z - w*y), 2*(y*z + w*x), 1 - 2*(x*x + y*y),
120+ ], dim=-1)
121+ return R.reshape(*quat.shape[:-1], 3, 3)
Mexamples/rasterization/config.py+119-105
@@ -1,105 +1,119 @@
1-# coding=utf-81+# coding=utf-8
2-# Adapted from2+# Adapted from
3-# https://github.com/nerfstudio-project/gsplat/blob/main/examples/simple_trainer.py3+# https://github.com/nerfstudio-project/gsplat/blob/main/examples/simple_trainer.py
4-# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.4+# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
5- 5+ 
6-from dataclasses import dataclass, field6+from dataclasses import dataclass, field
7-from typing import List, Optional7+from typing import List, Optional
8-from typing_extensions import Literal8+from typing_extensions import Literal
9- 9+ 
10- 10+ 
11-@dataclass11+@dataclass
12-class Config:12+class Config:
13- # Path to the .pt files. If provide, it will skip training and run evaluation only.13+ # Path to the .pt files. If provide, it will skip training and run evaluation only.
14- ckpt: Optional[List[str]] = None14+ ckpt: Optional[List[str]] = None
15- 15+ 
16- # Path to the Mip-NeRF 360 dataset16+ # Path to the Mip-NeRF 360 dataset
17- data_dir: str = "data/360_v2/garden"17+ data_dir: str = "data/360_v2/garden"
18- # Downsample factor for the dataset18+ # Downsample factor for the dataset
19- data_factor: int = 419+ data_factor: int = 4
20- # Directory to save results20+ # Directory to save results
21- result_dir: str = "results/garden"21+ result_dir: str = "results/garden"
22- 22+ 
23- # Every N images there is a test image23+ # Every N images there is a test image
24- test_every: int = 824+ test_every: int = 8
25- # Random crop size for training (experimental)25+ # Random crop size for training (experimental)
26- patch_size: Optional[int] = None26+ patch_size: Optional[int] = None
27- # A global scaler that applies to the scene size related parameters27+ # A global scaler that applies to the scene size related parameters
28- global_scale: float = 1.028+ global_scale: float = 1.0
29- # Normalize the world space29+ # Normalize the world space
30- normalize_world_space: bool = True30+ normalize_world_space: bool = True
31- # Camera model31+ # Camera model
32- camera_model: Literal["pinhole", "ortho", "fisheye"] = "pinhole"32+ camera_model: Literal["pinhole", "ortho", "fisheye"] = "pinhole"
33- 33+ 
34- # Port for the viewer server34+ # Port for the viewer server
35- port: int = 808035+ port: int = 8080
36- 36+ 
37- # tile size: 32 or 6437+ # tile size: 32 or 64
38- tile_size: int = 3238+ tile_size: int = 32
39- # Batch size for training. Learning rates are scaled automatically39+ # Batch size for training. Learning rates are scaled automatically
40- batch_size: int = 140+ batch_size: int = 1
41- # A global factor to scale the number of training steps41+ # A global factor to scale the number of training steps
42- steps_scaler: float = 1.042+ steps_scaler: float = 1.0
43- 43+ 
44- # Number of training steps44+ # Number of training steps
45- max_steps: int = 30_00045+ max_steps: int = 30_000
46- # Steps to evaluate the model46+ # Steps to evaluate the model
47- eval_steps: List[int] = field(default_factory=lambda: [500, 7_000, 30_000])47+ eval_steps: List[int] = field(default_factory=lambda: [500, 7_000, 30_000])
48- # Steps to save the model48+ # Steps to save the model
49- save_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])49+ save_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])
50- # Whether to save ply file (storage size can be large)50+ # Whether to save ply file (storage size can be large)
51- save_ply: bool = False51+ save_ply: bool = False
52- # Steps to save the model as ply52+ # Steps to save the model as ply
53- ply_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])53+ ply_steps: List[int] = field(default_factory=lambda: [7_000, 30_000])
54- 54+ 
55- # Initialization strategy55+ # Initialization strategy
56- init_type: str = "sfm"56+ init_type: str = "sfm"
57- # Initial number of GSs. Ignored if using sfm57+ # Initial number of GSs. Ignored if using sfm
58- init_num_pts: int = 100_00058+ init_num_pts: int = 100_000
59- # Initial extent of GSs as a multiple of the camera extent. Ignored if using sfm59+ # Initial extent of GSs as a multiple of the camera extent. Ignored if using sfm
60- init_extent: float = 3.060+ init_extent: float = 3.0
61- # Degree of spherical harmonics61+ # Degree of spherical harmonics
62- sh_degree: int = 362+ sh_degree: int = 3
63- # Turn on another SH degree every this steps63+ # Turn on another SH degree every this steps
64- sh_degree_interval: int = 100064+ sh_degree_interval: int = 1000
65- # Initial opacity of GS65+ # Initial opacity of GS
66- init_opa: float = 0.166+ init_opa: float = 0.1
67- # Initial scale of GS67+ # Initial scale of GS
68- init_scale: float = 1.068+ init_scale: float = 1.0
69- # Weight for SSIM loss69+ # Weight for SSIM loss
70- ssim_lambda: float = 0.270+ ssim_lambda: float = 0.2
71- 71+ 
72- # Near plane clipping distance72+ # Near plane clipping distance
73- near_plane: float = 0.0173+ near_plane: float = 0.01
74- # Far plane clipping distance74+ # Far plane clipping distance
75- far_plane: float = 1e1075+ far_plane: float = 1e10
76- 76+ 
77- # Strategy for GS densification77+ # Strategy for GS densification
78- strategy: Optional = None78+ strategy: Optional = None
79- 79+ 
80- # Use sparse gradients for optimization. (experimental)80+ # Use sparse gradients for optimization. (experimental)
81- sparse_grad: bool = False81+ sparse_grad: bool = False
82- 82+ 
83- # Opacity regularization83+ # Opacity regularization
84- opacity_reg: float = 0.084+ opacity_reg: float = 0.0
85- # Scale regularization85+ # Scale regularization
86- scale_reg: float = 0.086+ scale_reg: float = 0.0
87- 87+ 
88- # Enable depth loss. (experimental)88+ # Enable depth loss. (experimental)
89- depth_loss: bool = False89+ depth_loss: bool = False
90- # Weight for depth loss90+ # Weight for depth loss
91- depth_lambda: float = 1e-291+ depth_lambda: float = 1e-2
92- 92+ 
93- # Dump information to tensorboard every this steps93+ # Dump information to tensorboard every this steps
94- tb_every: int = 10094+ tb_every: int = 100
95- # Save training images to tensorboard95+ # Save training images to tensorboard
96- tb_save_image: bool = False96+ tb_save_image: bool = False
97- 97+ 
98- lpips_net: Literal["vgg", "alex"] = "alex"98+ lpips_net: Literal["vgg", "alex"] = "alex"
99- 99+
100- def adjust_steps(self, factor: float):100+ 
101- self.eval_steps = [int(i * factor) for i in self.eval_steps]101+ # === Densification (gsplat DefaultStrategy) ===
102- self.save_steps = [int(i * factor) for i in self.save_steps]102+ densify: bool = False
103- self.ply_steps = [int(i * factor) for i in self.ply_steps]103+ prune_opa: float = 0.005
104- self.max_steps = int(self.max_steps * factor)104+ grow_grad2d: float = 0.0002
105- self.sh_degree_interval = int(self.sh_degree_interval * factor)105+ refine_start_iter: int = 500
106+ refine_stop_iter: int = 15000
107+ refine_every: int = 100
108+ reset_every: int = 3000
109+ 
110+ def adjust_steps(self, factor: float):
111+ self.eval_steps = [int(i * factor) for i in self.eval_steps]
112+ self.save_steps = [int(i * factor) for i in self.save_steps]
113+ self.ply_steps = [int(i * factor) for i in self.ply_steps]
114+ self.max_steps = int(self.max_steps * factor)
115+ self.sh_degree_interval = int(self.sh_degree_interval * factor)
116+ self.refine_start_iter = int(self.refine_start_iter * factor)
117+ self.refine_stop_iter = int(self.refine_stop_iter * factor)
118+ self.refine_every = int(self.refine_every * factor)
119+ self.reset_every = int(self.reset_every * factor)
Mexamples/rasterization/runner.py+53-1
@@ -26,7 +26,22 @@ from rasterization.config import Config
26 26 
27from gauss_splat import Rasterizer as _GSRasterizer27from gauss_splat import Rasterizer as _GSRasterizer
28 28 
29-class Rasterizer:29+# === NPU densification support ===
30+from gsplat.strategy import DefaultStrategy
31+ 
32+ 
33+def _to_CN2(t):
34+ """Convert means2d gradient / radii to [C, N, 2] format."""
35+ while t.ndim > 3 and t.shape[0] == 1:
36+ t = t.squeeze(0)
37+ if t.ndim == 3 and t.shape[1] in (2, 3) and t.shape[1] < t.shape[2]:
38+ t = t.permute(0, 2, 1).contiguous()
39+ elif t.ndim == 2:
40+ t = t.unsqueeze(0)
41+ return t
atomgit-bot
atomgit-botatomgit-bot29 天前

🟠 High Priority

_to_CN2 函数被同时用于处理 means2d 梯度（4D tensor [1, C, N, 2]）和 radii（3D tensor [1, C, N]），但该函数只能正确处理 4D→3D 的转换：

在 DefaultStrategy._update_state() 中，info["radii"] 被用于布尔索引 sel = (info["radii"] > 0.0).all(dim=-1)，然后 grads[sel] 对 [C, N, 2] 的 grads 做索引。当 radii 形状不正确时（如 C=1 时 [1, 1, N] → sel 形状 [1,1]），PyTorch 的布尔索引会将 2D mask 广播到 grads 的前两维，导致 grads_sel 要么只筛选出 1 个元素，要么返回全部 2*N_vis 个元素作为 1D tensor，截断后梯度归因完全错误。这会使 densification 的 grad2d 累积统计失效，导致错误的 grow/prune 决策。

建议：为 radii 单独提供形状转换：新增 _to_CN() 函数，只 squeeze 前导 batch 维度（while ndim > 2 and shape[0] == 1），在第 343 行将 _to_CN2(info["radii"]) 替换为 _to_CN(info["radii"])。同时保留 _to_CN2 仅用于 means2d 梯度。

改动建议
41
+ def _to_CN2(t):
42
+ """Convert means2d gradient to [C, N, 2] format."""
43
+ while t.ndim > 3 and t.shape[0] == 1:
44
+ t = t.squeeze(0)
45
+ if t.ndim == 3 and t.shape[1] in (2, 3) and t.shape[1] < t.shape[2]:
46
+ t = t.permute(0, 2, 1).contiguous()
47
+ elif t.ndim == 2:
48
+ t = t.unsqueeze(0)
41
- return t
49
+ return t
50
+
51
+
52
+ def _to_CN(t):
53
+ """Convert radii-like tensor to [C, N] format."""
54
+ while t.ndim > 2 and t.shape[0] == 1:
55
+ t = t.squeeze(0)
56
+ if t.ndim == 1:
57
+ t = t.unsqueeze(0)
58
+ return t
应用建议
likedislike
42+ 
43+ 
44+class Rasterizer:
30 def __init__(self, cfg: Config) -> None:45 def __init__(self, cfg: Config) -> None:
31 self.cfg = cfg46 self.cfg = cfg
32 self._impl = _GSRasterizer()47 self._impl = _GSRasterizer()
@@ -119,6 +134,25 @@ class Runner:
119 world_size=world_size,134 world_size=world_size,
120 )135 )
121 print("Model initialized. Number of GS:", len(self.splats["means"]))136 print("Model initialized. Number of GS:", len(self.splats["means"]))
137+ if getattr(cfg, "densify", False):
138+ self.strategy = DefaultStrategy(
139+ prune_opa=cfg.prune_opa,
140+ grow_grad2d=cfg.grow_grad2d,
141+ refine_start_iter=cfg.refine_start_iter,
142+ refine_stop_iter=cfg.refine_stop_iter,
143+ refine_every=cfg.refine_every,
144+ reset_every=cfg.reset_every,
145+ verbose=(world_rank == 0),
146+ )
147+ self.strategy.check_sanity(self.splats, self.optimizers)
148+ self.strategy_state = self.strategy.initialize_state(scene_scale=self.scene_scale)
149+ if world_rank == 0:
150+ print(f"[Densify] DefaultStrategy enabled "
151+ f"(refine [{cfg.refine_start_iter}-{cfg.refine_stop_iter}] every {cfg.refine_every}, "
152+ f"reset every {cfg.reset_every}, grow_grad2d={cfg.grow_grad2d})")
153+ else:
154+ self.strategy = None
155+ self.strategy_state = None
122 156 
123 # Losses & Metrics.157 # Losses & Metrics.
124 self.ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)158 self.ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(self.device)
@@ -241,6 +275,12 @@ class Runner:
241 + cfg.scale_reg * torch.abs(torch.exp(self.splats["scales"])).mean()275 + cfg.scale_reg * torch.abs(torch.exp(self.splats["scales"])).mean()
242 )276 )
243 277 
278+ _orig_m2d = None
279+ if self.strategy is not None:
280+ _orig_m2d = info["means2d"]
281+ if _orig_m2d.requires_grad:
282+ _orig_m2d.retain_grad()
283+ 
244 loss.backward()284 loss.backward()
245 285 
246 desc = f"loss={loss.item():.3f}| " f"sh degree={sh_degree_to_use}| "286 desc = f"loss={loss.item():.3f}| " f"sh degree={sh_degree_to_use}| "
@@ -299,6 +339,18 @@ class Runner:
299 for scheduler in schedulers:339 for scheduler in schedulers:
300 scheduler.step()340 scheduler.step()
301 341 
342+ if self.strategy is not None and _orig_m2d is not None and _orig_m2d.grad is not None:
343+ with torch.no_grad():
344+ g = _to_CN2(_orig_m2d.grad)
345+ radii_cn2 = _to_CN2(info["radii"])
346+ proxy_m2d = torch.empty_like(g, requires_grad=False)
347+ proxy_m2d.grad = g
348+ proxy_info = dict(info)
349+ proxy_info["means2d"] = proxy_m2d
350+ proxy_info["radii"] = radii_cn2
351+ self.strategy.step_post_backward(self.splats, self.optimizers,
352+ self.strategy_state, step, proxy_info)
353+ 
302 # eval the full set354 # eval the full set
303 if step in [i - 1 for i in cfg.eval_steps]:355 if step in [i - 1 for i in cfg.eval_steps]:
304 self.eval(step)356 self.eval(step)
Mgauss_splat/ops/projection_three_dims_gaussian_fused.py+2-2
@@ -93,7 +93,7 @@ class ProjectionThreeDimsGaussianFused(Function):
93 ctx.width = width93 ctx.width = width
94 ctx.height = height94 ctx.height = height
95 return means2d_culling, depths_culling, conics_culling, opacities_culling,\95 return means2d_culling, depths_culling, conics_culling, opacities_culling,\
96- radius_culling, covars2d_culling, colors_culling, cnt96+ radius_culling, covars2d_culling, colors_culling, cnt, proj_filter
97 97 
98 @staticmethod98 @staticmethod
99 # pylint: disable=too-many-arguments,huawei-too-many-arguments,too-many-return-values99 # pylint: disable=too-many-arguments,huawei-too-many-arguments,too-many-return-values
@@ -103,7 +103,7 @@ class ProjectionThreeDimsGaussianFused(Function):
103 means, conics, viewmats, quats, scales, ks, proj_filter, compensations = ctx.saved_tensors103 means, conics, viewmats, quats, scales, ks, proj_filter, compensations = ctx.saved_tensors
104 width = ctx.width104 width = ctx.width
105 height = ctx.height105 height = ctx.height
106- v_means2d, v_depths, v_conics, v_opacities_culling, v_radii, v_covar2d, v_colors_culling, v_cnt = v_args106+ v_means2d, v_depths, v_conics, v_opacities_culling, v_radii, v_covar2d, v_colors_culling, v_cnt, _v_proj_filter = v_args
107 v_pw, v_quats, v_scales, v_r, v_colors, v_opacities = fully_fused_projection_bwd(107 v_pw, v_quats, v_scales, v_r, v_colors, v_opacities = fully_fused_projection_bwd(
108 means,108 means,
109 quats,109 quats,
Mgauss_splat/ops/rendering.py+3-1
@@ -165,7 +165,7 @@ class Rasterizer:
165 colors = spherical_harmonics(sh_degree, rays_d.reshape(B, N, 3), shs[0, :, :k, :].reshape(B, N, k, 3))165 colors = spherical_harmonics(sh_degree, rays_d.reshape(B, N, 3), shs[0, :, :k, :].reshape(B, N, k, 3))
166 colors = (colors + 0.5).clip(min=0.0)166 colors = (colors + 0.5).clip(min=0.0)
167 167 
168- means2d, depths, conics, opacities, radius, covars2d, colors, cnt = projection_three_dims_gaussian_fused(168+ means2d, depths, conics, opacities, radius, covars2d, colors, cnt, proj_filter = projection_three_dims_gaussian_fused(
169 means.reshape(B, N, 3),169 means.reshape(B, N, 3),
170 colors,170 colors,
171 None,171 None,
@@ -267,5 +267,7 @@ class Rasterizer:
267 "width": width,267 "width": width,
268 "height": height,268 "height": height,
269 "n_cameras": C,269 "n_cameras": C,
270+ "proj_filter": proj_filter,
271+ "n_gaussians": N,
270 }272 }
271 return render_colors, render_depths, meta273 return render_colors, render_depths, meta
Mtests/test_projection_3dgs_forward.py+1-1
@@ -562,7 +562,7 @@ class TestProjection3DGSForward(TestCase):
562 viewmats = inputs.viewmats.npu()562 viewmats = inputs.viewmats.npu()
563 ks = inputs.ks.npu()563 ks = inputs.ks.npu()
564 means2d_culling, depths_culling, conics_culling, opacities_culling, \564 means2d_culling, depths_culling, conics_culling, opacities_culling, \
565- radius_culling, covars2d_culling, colors_culling, cnt = \565+ radius_culling, covars2d_culling, colors_culling, cnt, _proj_filter = \
566 projection_three_dims_gaussian_fused(means, colors, covars, \566 projection_three_dims_gaussian_fused(means, colors, covars, \
567 None, None, opacities, viewmats, ks, width, height, eps2d, \567 None, None, opacities, viewmats, ks, width, height, eps2d, \
568 near_plane, far_plane, calc_compensations, camera_model)568 near_plane, far_plane, calc_compensations, camera_model)
Mtests/test_rasterization.py+1-1
@@ -93,7 +93,7 @@ def _golden_rasterization(splats, cam, size, tile_size, active_sh_degree, camera
93 colors = (colors + 0.5).clip(min=0.0)93 colors = (colors + 0.5).clip(min=0.0)
94 94 
95 from gauss_splat import projection_three_dims_gaussian_fused95 from gauss_splat import projection_three_dims_gaussian_fused
96- means2d, depths, conics, opacities_proj, radius, covars2d, colors_proj, cnt = \96+ means2d, depths, conics, opacities_proj, radius, covars2d, colors_proj, cnt, _proj_filter = \
97 projection_three_dims_gaussian_fused(97 projection_three_dims_gaussian_fused(
98 means.reshape(B, N, 3), colors, None,98 means.reshape(B, N, 3), colors, None,
99 quats.reshape(B, N, 4), scales.reshape(B, N, 3), opacities.reshape(B, N),99 quats.reshape(B, N, 4), scales.reshape(B, N, 3), opacities.reshape(B, N),