Skip to content

Commit 4ed380f

Browse files
fwilliamsfwilliams
authored
Expose k and t parameters in MCMC relocation (#402)
These parameters control the falloff curve for noise added to each Gaussian during optimization as a function of scale and opacity (more opaque Gaussians are considered more certain and undergo smaller perturbation during optimization). This PR doesn't change the default values of these but exposes them so they can be tuned as hyperparameters. --------- Signed-off-by: Francis Williams <francis@fwilliams.info>
1 parent 1428cce commit 4ed380f

File tree

8 files changed

+53
-26
lines changed

8 files changed

+53
-26
lines changed

‎fvdb/_fvdb_cpp.pyi‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -325,7 +325,7 @@ class GaussianSplat3d:
325325
n_max: int,
326326
min_opacity: float,
327327
) -> tuple[torch.Tensor, torch.Tensor]: ...
328-
def add_noise_to_means(self, noise_scale: float) -> None: ...
328+
def add_noise_to_means(self, noise_scale: float, t: float = ..., k: float = ...) -> None: ...
329329
def reset_accumulated_gradient_state(self) -> None: ...
330330
def save_ply(self, filename: str, metadata: dict[str, str | int | float | torch.Tensor] | None) -> None: ...
331331
@staticmethod

‎fvdb/gaussian_splatting.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2592,14 +2592,16 @@ def relocate_gaussians(
25922592
min_opacity,
25932593
)
25942594

2595-
def add_noise_to_means(self, noise_scale: float) -> None:
2595+
def add_noise_to_means(self, noise_scale: float, t: float = 0.005, k: float = 100.0) -> None:
25962596
"""
25972597
Add noise to the Gaussian positions (means), scaled by ``noise_scale``.
25982598
25992599
Args:
26002600
noise_scale (float): Noise scale factor applied to scale-dependent noise.
2601+
t (float): Parameter t for noise scaling. Defaults to 0.005.
2602+
k (float): Parameter k for noise scaling. Defaults to 100.0.
26012603
"""
2602-
self._impl.add_noise_to_means(noise_scale)
2604+
self._impl.add_noise_to_means(noise_scale, t, k)
26032605

26042606
def reset_accumulated_gradient_state(self) -> None:
26052607
"""

‎src/fvdb/GaussianSplat3d.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1066,10 +1066,10 @@ GaussianSplat3d::relocateGaussians(const torch::Tensor &logScales,
10661066
}
10671067

10681068
void
1069-
GaussianSplat3d::addNoiseToMeans(const float noiseScale) {
1069+
GaussianSplat3d::addNoiseToMeans(const float noiseScale, const float t, const float k) {
10701070
FVDB_DISPATCH_KERNEL(mMeans.device(), [&]() {
10711071
return detail::ops::dispatchGaussianMCMCAddNoise<DeviceTag>(
1072-
mMeans, mLogScales, mLogitOpacities, mQuats, noiseScale);
1072+
mMeans, mLogScales, mLogitOpacities, mQuats, noiseScale, t, k);
10731073
});
10741074
}
10751075

‎src/fvdb/GaussianSplat3d.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1207,7 +1207,9 @@ class GaussianSplat3d {
12071207

12081208
/// @brief Add noise to the Gaussian positions (means), scaled by noiseScale.
12091209
/// @param noiseScale Noise scale
1210-
void addNoiseToMeans(const float noiseScale);
1210+
/// @param t Cutoff for opacity scaling
1211+
/// @param k Exponent for opacity scaling
1212+
void addNoiseToMeans(const float noiseScale, const float t = 0.005, const float k = 100.0);
12111213

12121214
/// @brief Select a subset of the Gaussians in this scene based on the given slice.
12131215
/// @param begin The start index of the slice (inclusive)

‎src/fvdb/detail/ops/gsplat/GaussianMCMCAddNoise.cu‎

Lines changed: 29 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ namespace fvdb::detail::ops {
1616

1717
template <typename ScalarType>
1818
inline __device__ ScalarType
19-
logistic(ScalarType x, ScalarType k = 100, ScalarType x0 = 0.995) {
20-
return 1 / (1 + exp(-k * (x - x0)));
19+
sigmoid(ScalarType x) {
20+
return ScalarType(1) / (ScalarType(1) + ::cuda::std::exp(-x));
2121
}
2222

2323
template <typename ScalarType>
@@ -27,23 +27,30 @@ gaussianMCMCAddNoiseKernel(fvdb::TorchRAcc64<ScalarType, 2> outMeans,
2727
fvdb::TorchRAcc64<ScalarType, 1> logitOpacities,
2828
fvdb::TorchRAcc64<ScalarType, 2> quats,
2929
fvdb::TorchRAcc64<ScalarType, 2> baseNoise,
30-
ScalarType noiseScale) {
30+
const ScalarType noiseScale,
31+
const ScalarType t,
32+
const ScalarType k) {
3133
const auto N = outMeans.size(0);
3234
for (uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < N;
3335
idx += blockDim.x * gridDim.x) {
34-
auto opacity = ScalarType(1.0) / (1 + exp(-logitOpacities[idx]));
36+
const auto opacity = sigmoid(logitOpacities[idx]);
3537

3638
const auto quatAcc = quats[idx];
3739
const auto logScaleAcc = logScales[idx];
38-
auto covar = quaternionAndScaleToCovariance<ScalarType>(
40+
const auto covar = quaternionAndScaleToCovariance<ScalarType>(
3941
nanovdb::math::Vec4<ScalarType>(quatAcc[0], quatAcc[1], quatAcc[2], quatAcc[3]),
4042
nanovdb::math::Vec3<ScalarType>(::cuda::std::exp(logScaleAcc[0]),
4143
::cuda::std::exp(logScaleAcc[1]),
4244
::cuda::std::exp(logScaleAcc[2])));
4345

4446
nanovdb::math::Vec3<ScalarType> noise = {
4547
baseNoise[idx][0], baseNoise[idx][1], baseNoise[idx][2]};
46-
noise *= logistic(1 - opacity) * noiseScale;
48+
49+
// The noise term is scaled down based on the opacity of the Gaussian.
50+
// More opaque Gaussians get less noise added to them.
51+
// The parameters t and k control the transition point and sharpness
52+
// of the scaling function.
53+
noise *= sigmoid(-k * (opacity - t)) * noiseScale;
4754
noise = covar * noise;
4855
outMeans[idx][0] += noise[0];
4956
outMeans[idx][1] += noise[1];
@@ -57,7 +64,9 @@ launchGaussianMCMCAddNoise(torch::Tensor &means, // [N, 3]
5764
const torch::Tensor &logScales, // [N, 3]
5865
const torch::Tensor &logitOpacities, // [N]
5966
const torch::Tensor &quats, // [N, 4]
60-
ScalarType noiseScale) {
67+
const ScalarType noiseScale,
68+
const ScalarType t,
69+
const ScalarType k) {
6170
const auto N = means.size(0);
6271

6372
const int blockDim = DEFAULT_BLOCK_DIM;
@@ -72,7 +81,9 @@ launchGaussianMCMCAddNoise(torch::Tensor &means, // [N, 3]
7281
logitOpacities.packed_accessor64<ScalarType, 1, torch::RestrictPtrTraits>(),
7382
quats.packed_accessor64<ScalarType, 2, torch::RestrictPtrTraits>(),
7483
baseNoise.packed_accessor64<ScalarType, 2, torch::RestrictPtrTraits>(),
75-
noiseScale);
84+
noiseScale,
85+
t,
86+
k);
7687

7788
C10_CUDA_KERNEL_LAUNCH_CHECK();
7889
}
@@ -83,13 +94,15 @@ dispatchGaussianMCMCAddNoise<torch::kCUDA>(torch::Tensor &means,
8394
const torch::Tensor &logScales, // [N]
8495
const torch::Tensor &logitOpacities, // [N]
8596
const torch::Tensor &quats, // [N, 4]
86-
float noiseScale) { // [N]
97+
const float noiseScale,
98+
const float t,
99+
const float k) {
87100
FVDB_FUNC_RANGE();
88101
const at::cuda::OptionalCUDAGuard device_guard(device_of(means));
89102

90103
const auto N = means.size(0);
91104

92-
launchGaussianMCMCAddNoise<float>(means, logScales, logitOpacities, quats, noiseScale);
105+
launchGaussianMCMCAddNoise<float>(means, logScales, logitOpacities, quats, noiseScale, t, k);
93106
}
94107

95108
template <>
@@ -98,7 +111,9 @@ dispatchGaussianMCMCAddNoise<torch::kPrivateUse1>(torch::Tensor &means, // [N, 3
98111
const torch::Tensor &logScales, // [N, 3]
99112
const torch::Tensor &logitOpacities, // [N]
100113
const torch::Tensor &quats, // [N, 4]
101-
float noiseScale) { // [N]
114+
const float noiseScale,
115+
const float t,
116+
const float k) {
102117
TORCH_CHECK_NOT_IMPLEMENTED(false, "GaussianMCMCAddNoise is not implemented for PrivateUse1");
103118
}
104119

@@ -108,7 +123,9 @@ dispatchGaussianMCMCAddNoise<torch::kCPU>(torch::Tensor &means, // [N,
108123
const torch::Tensor &logScales, // [N, 3]
109124
const torch::Tensor &logitOpacities, // [N]
110125
const torch::Tensor &quats, // [N, 4]
111-
float noiseScale) { // [N]
126+
const float noiseScale,
127+
const float t,
128+
const float k) {
112129
TORCH_CHECK_NOT_IMPLEMENTED(false, "GaussianMCMCAddNoise is not implemented for CPU");
113130
}
114131

‎src/fvdb/detail/ops/gsplat/GaussianMCMCAddNoise.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@ void dispatchGaussianMCMCAddNoise(torch::Tensor &means, // [N, 3]
1616
const torch::Tensor &logScales, // [N]
1717
const torch::Tensor &logitOpacities, // [N]
1818
const torch::Tensor &quats, // [N, 4]
19-
float noiseScale); // [N]
19+
const float noiseScale,
20+
const float t,
21+
const float k);
2022

2123
} // namespace ops
2224
} // namespace detail

‎src/python/GaussianSplatBinding.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -347,7 +347,11 @@ bind_gaussian_splat3d(py::module &m) {
347347
py::arg("n_max"),
348348
py::arg("min_opacity"))
349349

350-
.def("add_noise_to_means", &fvdb::GaussianSplat3d::addNoiseToMeans, py::arg("noise_scale"))
350+
.def("add_noise_to_means",
351+
&fvdb::GaussianSplat3d::addNoiseToMeans,
352+
py::arg("noise_scale"),
353+
py::arg("t") = 0.005,
354+
py::arg("k") = 100.0)
351355

352356
.def("index_select", &fvdb::GaussianSplat3d::indexSelect, py::arg("indices"))
353357
.def("mask_select", &fvdb::GaussianSplat3d::maskSelect, py::arg("mask"))

‎src/tests/GaussianMCMCAddNoiseTest.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ TEST_F(GaussianMCMCAddNoiseTest, AppliesNoiseWithDeterministicBaseNoise) {
6464
auto meansBaseline = means.clone();
6565

6666
fvdb::detail::ops::dispatchGaussianMCMCAddNoise<torch::kCUDA>(
67-
means, logScales, logitOpacities, quats, noiseScale);
67+
means, logScales, logitOpacities, quats, noiseScale, 0.005, 100.0);
6868

6969
restoreCudaGeneratorState(rngState);
7070
const auto baseNoise = torch::randn_like(meansBaseline);
@@ -91,7 +91,7 @@ TEST_F(GaussianMCMCAddNoiseTest, RespectsAnisotropicScales) {
9191
const auto rngState = saveCudaGeneratorState();
9292

9393
fvdb::detail::ops::dispatchGaussianMCMCAddNoise<torch::kCUDA>(
94-
means, logScales, logitOpacities, quats, noiseScale);
94+
means, logScales, logitOpacities, quats, noiseScale, 0.005, 100);
9595

9696
restoreCudaGeneratorState(rngState);
9797
const auto baseNoise = torch::randn_like(means);
@@ -115,7 +115,7 @@ TEST_F(GaussianMCMCAddNoiseTest, HighOpacitySuppressesNoise) {
115115
constexpr float noiseScale = 1.0f;
116116

117117
fvdb::detail::ops::dispatchGaussianMCMCAddNoise<torch::kCUDA>(
118-
means, logScales, logitOpacities, quats, noiseScale);
118+
means, logScales, logitOpacities, quats, noiseScale, 0.005, 100.0);
119119

120120
// Gate approaches zero when opacity ~1; expect negligible movement.
121121
const auto maxAbs = torch::abs(means).max().item<float>();
@@ -133,7 +133,7 @@ TEST_F(GaussianMCMCAddNoiseTest, ZeroNoiseScaleNoOp) {
133133
floatOpts());
134134

135135
fvdb::detail::ops::dispatchGaussianMCMCAddNoise<torch::kCUDA>(
136-
means, logScales, logitOpacities, quats, /*noiseScale=*/0.0f);
136+
means, logScales, logitOpacities, quats, /*noiseScale=*/0.0f, 0.005, 100.0);
137137

138138
EXPECT_TRUE(torch::allclose(means, origMeans));
139139
}
@@ -146,15 +146,15 @@ TEST_F(GaussianMCMCAddNoiseTest, CpuAndPrivateUseNotImplemented) {
146146
torch::tensor({{1.0f, 0.0f, 0.0f, 0.0f}}, fvdb::test::tensorOpts<float>(torch::kCPU));
147147

148148
EXPECT_THROW((fvdb::detail::ops::dispatchGaussianMCMCAddNoise<torch::kCPU>(
149-
means, logScales, logitOpacities, quats, 1.0f)),
149+
means, logScales, logitOpacities, quats, 1.0f, 0.005, 100)),
150150
c10::Error);
151151

152152
auto meansCuda = means.cuda();
153153
auto logScalesCuda = logScales.cuda();
154154
auto logitOpacitiesCuda = logitOpacities.cuda();
155155
auto quatsCuda = quats.cuda();
156156
EXPECT_THROW((fvdb::detail::ops::dispatchGaussianMCMCAddNoise<torch::kPrivateUse1>(
157-
meansCuda, logScalesCuda, logitOpacitiesCuda, quatsCuda, 1.0f)),
157+
meansCuda, logScalesCuda, logitOpacitiesCuda, quatsCuda, 1.0f, 0.005, 100.0)),
158158
c10::Error);
159159
}
160160

0 commit comments

Comments
 (0)