mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-10 18:17:31 +08:00
port gsplat impl over to splatfacto
This commit is contained in:
Justin Kerr
parent
f97eb2e5f0
commit
bfdda8ca6a
@ -44,6 +44,7 @@ from nerfstudio.models.base_model import Model, ModelConfig
|
||||
from nerfstudio.utils.colors import get_color
|
||||
from nerfstudio.utils.misc import torch_compile
|
||||
from nerfstudio.utils.rich_utils import CONSOLE
|
||||
from gsplat.relocation import compute_relocation
|
||||
|
||||
|
||||
def random_quat_tensor(N):
|
||||
@ -350,49 +351,34 @@ class SplatfactoModel(Model):
|
||||
# Exclude the point itself from the result and return
|
||||
return distances[:, 1:].astype(np.float32), indices[:, 1:].astype(np.float32)
|
||||
|
||||
def remove_from_optim(self, optimizer, deleted_mask, new_params):
|
||||
"""removes the deleted_mask from the optimizer provided"""
|
||||
assert len(new_params) == 1
|
||||
# assert isinstance(optimizer, torch.optim.Adam), "Only works with Adam"
|
||||
|
||||
param = optimizer.param_groups[0]["params"][0]
|
||||
param_state = optimizer.state[param]
|
||||
del optimizer.state[param]
|
||||
|
||||
# Modify the state directly without deleting and reassigning.
|
||||
if "exp_avg" in param_state:
|
||||
param_state["exp_avg"] = param_state["exp_avg"][~deleted_mask]
|
||||
param_state["exp_avg_sq"] = param_state["exp_avg_sq"][~deleted_mask]
|
||||
|
||||
# Update the parameter in the optimizer's param group.
|
||||
del optimizer.param_groups[0]["params"][0]
|
||||
del optimizer.param_groups[0]["params"]
|
||||
optimizer.param_groups[0]["params"] = new_params
|
||||
optimizer.state[new_params[0]] = param_state
|
||||
|
||||
def remove_from_all_optim(self, optimizers, deleted_mask):
|
||||
param_groups = self.get_gaussian_param_groups()
|
||||
for group, param in param_groups.items():
|
||||
self.remove_from_optim(optimizers.optimizers[group], deleted_mask, param)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def dup_in_optim(self, optimizer, dup_mask, new_params, n=2):
|
||||
def zero_in_optim(self, optimizer, zero_mask, new_params):
|
||||
"""adds the parameters to the optimizer"""
|
||||
param = optimizer.param_groups[0]["params"][0]
|
||||
param_state = optimizer.state[param]
|
||||
if "exp_avg" in param_state:
|
||||
param_state["exp_avg"][zero_mask] = 0
|
||||
param_state["exp_avg_sq"][zero_mask] = 0
|
||||
del optimizer.state[param]
|
||||
optimizer.state[new_params[0]] = param_state
|
||||
optimizer.param_groups[0]["params"] = new_params
|
||||
del param
|
||||
|
||||
def add_in_optim(self, optimizer, new_params, n_new):
|
||||
"""adds the parameters to the optimizer"""
|
||||
param = optimizer.param_groups[0]["params"][0]
|
||||
param_state = optimizer.state[param]
|
||||
if "exp_avg" in param_state:
|
||||
repeat_dims = (n,) + tuple(1 for _ in range(param_state["exp_avg"].dim() - 1))
|
||||
param_state["exp_avg"] = torch.cat(
|
||||
[
|
||||
param_state["exp_avg"],
|
||||
torch.zeros_like(param_state["exp_avg"][dup_mask.squeeze()]).repeat(*repeat_dims),
|
||||
torch.zeros_like(param_state["exp_avg"][0:1,...]).repeat(n_new, *[1 for _ in range(len(param_state["exp_avg"].shape)-1)]),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
param_state["exp_avg_sq"] = torch.cat(
|
||||
[
|
||||
param_state["exp_avg_sq"],
|
||||
torch.zeros_like(param_state["exp_avg_sq"][dup_mask.squeeze()]).repeat(*repeat_dims),
|
||||
torch.zeros_like(param_state["exp_avg_sq"][0:1,...]).repeat(n_new, *[1 for _ in range(len(param_state["exp_avg"].shape)-1)]),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
@ -401,35 +387,21 @@ class SplatfactoModel(Model):
|
||||
optimizer.param_groups[0]["params"] = new_params
|
||||
del param
|
||||
|
||||
def dup_in_all_optim(self, optimizers, dup_mask, n):
|
||||
def zero_in_all_optim(self, optimizers, zero_mask):
|
||||
param_groups = self.get_gaussian_param_groups()
|
||||
for group, param in param_groups.items():
|
||||
self.dup_in_optim(optimizers.optimizers[group], dup_mask, param, n)
|
||||
self.zero_in_optim(optimizers.optimizers[group], zero_mask, param)
|
||||
|
||||
def add_in_all_optim(self, optimizers, n_new):
|
||||
param_groups = self.get_gaussian_param_groups()
|
||||
for group, param in param_groups.items():
|
||||
self.add_in_optim(optimizers.optimizers[group], param, n_new)
|
||||
|
||||
def after_train(self, step: int):
|
||||
assert step == self.step
|
||||
# to save some training time, we no longer need to update those stats post refinement
|
||||
if self.step >= self.config.stop_split_at:
|
||||
return
|
||||
with torch.no_grad():
|
||||
# keep track of a moving average of grad norms
|
||||
visible_mask = (self.radii > 0).flatten()
|
||||
grads = self.xys.absgrad[0][visible_mask].norm(dim=-1) # type: ignore
|
||||
# print(f"grad norm min {grads.min().item()} max {grads.max().item()} mean {grads.mean().item()} size {grads.shape}")
|
||||
if self.xys_grad_norm is None:
|
||||
self.xys_grad_norm = torch.zeros(self.num_points, device=self.device, dtype=torch.float32)
|
||||
self.vis_counts = torch.ones(self.num_points, device=self.device, dtype=torch.float32)
|
||||
assert self.vis_counts is not None
|
||||
self.vis_counts[visible_mask] += 1
|
||||
self.xys_grad_norm[visible_mask] += grads
|
||||
# update the max screen size, as a ratio of number of pixels
|
||||
if self.max_2Dsize is None:
|
||||
self.max_2Dsize = torch.zeros_like(self.radii, dtype=torch.float32)
|
||||
newradii = self.radii.detach()[visible_mask]
|
||||
self.max_2Dsize[visible_mask] = torch.maximum(
|
||||
self.max_2Dsize[visible_mask],
|
||||
newradii / float(max(self.last_size[0], self.last_size[1])),
|
||||
)
|
||||
|
||||
def set_crop(self, crop_box: Optional[OrientedBox]):
|
||||
self.crop_box = crop_box
|
||||
@ -437,175 +409,104 @@ class SplatfactoModel(Model):
|
||||
def set_background(self, background_color: torch.Tensor):
|
||||
assert background_color.shape == (3,)
|
||||
self.background_color = background_color
|
||||
|
||||
@torch.no_grad()
|
||||
def relocate_gs(self, optimizers, min_opacity: float = 0.005) -> int:
|
||||
dead_mask = torch.sigmoid(self.opacities) <= min_opacity
|
||||
dead_indices = dead_mask.nonzero(as_tuple=True)[0]
|
||||
alive_indices = (~dead_mask).nonzero(as_tuple=True)[0]
|
||||
num_gs = len(dead_indices)
|
||||
if num_gs <= 0:
|
||||
return num_gs
|
||||
|
||||
# Sample for new GSs
|
||||
eps = torch.finfo(torch.float32).eps
|
||||
probs = torch.sigmoid(self.opacities)[alive_indices].squeeze()
|
||||
probs = probs / (probs.sum() + eps)
|
||||
sampled_idxs = torch.multinomial(probs, num_gs, replacement=True)
|
||||
sampled_idxs = alive_indices[sampled_idxs]
|
||||
new_opacities, new_scales = compute_relocation(
|
||||
opacities=torch.sigmoid(self.opacities)[sampled_idxs].squeeze(),
|
||||
scales=torch.exp(self.scales)[sampled_idxs],
|
||||
ratios=torch.bincount(sampled_idxs)[sampled_idxs] + 1,
|
||||
)
|
||||
new_opacities = torch.clamp(new_opacities, max=1.0 - eps, min=min_opacity)
|
||||
self.gauss_params["opacities"][sampled_idxs] = torch.logit(new_opacities).unsqueeze(-1)
|
||||
self.gauss_params["scales"][sampled_idxs] = torch.log(new_scales)
|
||||
|
||||
# Update splats and optimizers
|
||||
for k in self.gauss_params.keys():
|
||||
self.gauss_params[k][dead_indices] = self.gauss_params[k][sampled_idxs]
|
||||
self.zero_in_all_optim(optimizers,sampled_idxs)
|
||||
torch.cuda.empty_cache()
|
||||
return num_gs
|
||||
|
||||
@torch.no_grad()
|
||||
def add_new_gs(self, cap_max: int, optimizers, min_opacity: float = 0.005) -> int:
|
||||
current_num_points = self.num_points
|
||||
target_num = min(cap_max, int(1.05 * current_num_points))
|
||||
num_gs = max(0, target_num - current_num_points)
|
||||
if num_gs <= 0:
|
||||
return num_gs
|
||||
|
||||
# Sample for new GSs
|
||||
eps = torch.finfo(torch.float32).eps
|
||||
probs = torch.sigmoid(self.opacities).squeeze()
|
||||
probs = probs / (probs.sum() + eps)
|
||||
sampled_idxs = torch.multinomial(probs, num_gs, replacement=True)
|
||||
new_opacities, new_scales = compute_relocation(
|
||||
opacities=torch.sigmoid(self.opacities)[sampled_idxs].squeeze(),
|
||||
scales=torch.exp(self.scales)[sampled_idxs],
|
||||
ratios=torch.bincount(sampled_idxs)[sampled_idxs] + 1,
|
||||
)
|
||||
new_opacities = torch.clamp(new_opacities, max=1.0 - eps, min=min_opacity)
|
||||
self.gauss_params["opacities"][sampled_idxs] = torch.logit(new_opacities).unsqueeze(-1)
|
||||
self.gauss_params["scales"][sampled_idxs] = torch.log(new_scales)
|
||||
|
||||
# Update splats and optimizers
|
||||
for k in self.gauss_params.keys():
|
||||
self.gauss_params[k] = torch.cat([self.gauss_params[k], self.gauss_params[k][sampled_idxs]])
|
||||
self.add_in_all_optim(optimizers, new_opacities.shape[0])
|
||||
torch.cuda.empty_cache()
|
||||
print(f"Added {num_gs} new gaussians, total now {self.num_points}")
|
||||
return num_gs
|
||||
|
||||
@torch.no_grad()
|
||||
def add_noise_to_gs(self, last_lr):
|
||||
opacities = torch.sigmoid(self.opacities)
|
||||
scales = torch.exp(self.scales)
|
||||
from gsplat import quat_scale_to_covar_preci
|
||||
|
||||
actual_covariance, _ = quat_scale_to_covar_preci(
|
||||
self.quats,
|
||||
scales,
|
||||
compute_covar=True,
|
||||
compute_preci=False,
|
||||
triu=False,
|
||||
)
|
||||
|
||||
def op_sigmoid(x, k=100, x0=0.995):
|
||||
return 1 / (1 + torch.exp(-k * (x - x0)))
|
||||
|
||||
noise = (
|
||||
torch.randn_like(self.means)
|
||||
* op_sigmoid(1 - opacities)
|
||||
* 5.0e5
|
||||
* last_lr
|
||||
)
|
||||
noise = torch.bmm(actual_covariance, noise.unsqueeze(-1)).squeeze(-1)
|
||||
self.gauss_params['means'] += noise
|
||||
|
||||
def refinement_after(self, optimizers: Optimizers, step):
|
||||
assert step == self.step
|
||||
if self.step <= self.config.warmup_length:
|
||||
return
|
||||
with torch.no_grad():
|
||||
# Offset all the opacity reset logic by refine_every so that we don't
|
||||
# save checkpoints right when the opacity is reset (saves every 2k)
|
||||
# then cull
|
||||
# only split/cull if we've seen every image since opacity reset
|
||||
reset_interval = self.config.reset_alpha_every * self.config.refine_every
|
||||
do_densification = (
|
||||
self.step < self.config.stop_split_at
|
||||
and self.step % reset_interval > self.num_train_data + self.config.refine_every
|
||||
)
|
||||
if do_densification:
|
||||
# then we densify
|
||||
assert self.xys_grad_norm is not None and self.vis_counts is not None and self.max_2Dsize is not None
|
||||
avg_grad_norm = (self.xys_grad_norm / self.vis_counts) * 0.5 * max(self.last_size[0], self.last_size[1])
|
||||
high_grads = (avg_grad_norm > self.config.densify_grad_thresh).squeeze()
|
||||
splits = (self.scales.exp().max(dim=-1).values > self.config.densify_size_thresh).squeeze()
|
||||
if self.step < self.config.stop_screen_size_at:
|
||||
splits |= (self.max_2Dsize > self.config.split_screen_size).squeeze()
|
||||
splits &= high_grads
|
||||
nsamps = self.config.n_split_samples
|
||||
split_params = self.split_gaussians(splits, nsamps)
|
||||
if self.step % self.config.refine_every == 0:
|
||||
last_lr = optimizers.optimizers["means"].param_groups[0]["lr"]
|
||||
self.relocate_gs(optimizers)
|
||||
self.add_new_gs(3_000_000, optimizers)
|
||||
self.add_noise_to_gs(last_lr)
|
||||
|
||||
dups = (self.scales.exp().max(dim=-1).values <= self.config.densify_size_thresh).squeeze()
|
||||
dups &= high_grads
|
||||
dup_params = self.dup_gaussians(dups)
|
||||
for name, param in self.gauss_params.items():
|
||||
self.gauss_params[name] = torch.nn.Parameter(
|
||||
torch.cat([param.detach(), split_params[name], dup_params[name]], dim=0)
|
||||
)
|
||||
# append zeros to the max_2Dsize tensor
|
||||
self.max_2Dsize = torch.cat(
|
||||
[
|
||||
self.max_2Dsize,
|
||||
torch.zeros_like(split_params["scales"][:, 0]),
|
||||
torch.zeros_like(dup_params["scales"][:, 0]),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
split_idcs = torch.where(splits)[0]
|
||||
self.dup_in_all_optim(optimizers, split_idcs, nsamps)
|
||||
|
||||
dup_idcs = torch.where(dups)[0]
|
||||
self.dup_in_all_optim(optimizers, dup_idcs, 1)
|
||||
|
||||
# After a guassian is split into two new gaussians, the original one should also be pruned.
|
||||
splits_mask = torch.cat(
|
||||
(
|
||||
splits,
|
||||
torch.zeros(
|
||||
nsamps * splits.sum() + dups.sum(),
|
||||
device=self.device,
|
||||
dtype=torch.bool,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
deleted_mask = self.cull_gaussians(splits_mask)
|
||||
elif self.step >= self.config.stop_split_at and self.config.continue_cull_post_densification:
|
||||
deleted_mask = self.cull_gaussians()
|
||||
else:
|
||||
# if we donot allow culling post refinement, no more gaussians will be pruned.
|
||||
deleted_mask = None
|
||||
|
||||
if deleted_mask is not None:
|
||||
self.remove_from_all_optim(optimizers, deleted_mask)
|
||||
|
||||
if self.step < self.config.stop_split_at and self.step % reset_interval == self.config.refine_every:
|
||||
# Reset value is set to be twice of the cull_alpha_thresh
|
||||
reset_value = self.config.cull_alpha_thresh * 2.0
|
||||
self.opacities.data = torch.clamp(
|
||||
self.opacities.data,
|
||||
max=torch.logit(torch.tensor(reset_value, device=self.device)).item(),
|
||||
)
|
||||
# reset the exp of optimizer
|
||||
optim = optimizers.optimizers["opacities"]
|
||||
param = optim.param_groups[0]["params"][0]
|
||||
param_state = optim.state[param]
|
||||
param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"])
|
||||
param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"])
|
||||
|
||||
self.xys_grad_norm = None
|
||||
self.vis_counts = None
|
||||
self.max_2Dsize = None
|
||||
|
||||
def cull_gaussians(self, extra_cull_mask: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
This function deletes gaussians with under a certain opacity threshold
|
||||
extra_cull_mask: a mask indicates extra gaussians to cull besides existing culling criterion
|
||||
"""
|
||||
n_bef = self.num_points
|
||||
# cull transparent ones
|
||||
culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze()
|
||||
below_alpha_count = torch.sum(culls).item()
|
||||
toobigs_count = 0
|
||||
if extra_cull_mask is not None:
|
||||
culls = culls | extra_cull_mask
|
||||
if self.step > self.config.refine_every * self.config.reset_alpha_every:
|
||||
# cull huge ones
|
||||
toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze()
|
||||
if self.step < self.config.stop_screen_size_at:
|
||||
# cull big screen space
|
||||
if self.max_2Dsize is not None:
|
||||
toobigs = toobigs | (self.max_2Dsize > self.config.cull_screen_size).squeeze()
|
||||
culls = culls | toobigs
|
||||
toobigs_count = torch.sum(toobigs).item()
|
||||
for name, param in self.gauss_params.items():
|
||||
self.gauss_params[name] = torch.nn.Parameter(param[~culls])
|
||||
|
||||
CONSOLE.log(
|
||||
f"Culled {n_bef - self.num_points} gaussians "
|
||||
f"({below_alpha_count} below alpha thresh, {toobigs_count} too bigs, {self.num_points} remaining)"
|
||||
)
|
||||
|
||||
return culls
|
||||
|
||||
def split_gaussians(self, split_mask, samps):
|
||||
"""
|
||||
This function splits gaussians that are too large
|
||||
"""
|
||||
n_splits = split_mask.sum().item()
|
||||
CONSOLE.log(f"Splitting {split_mask.sum().item()/self.num_points} gaussians: {n_splits}/{self.num_points}")
|
||||
centered_samples = torch.randn((samps * n_splits, 3), device=self.device) # Nx3 of axis-aligned scales
|
||||
scaled_samples = (
|
||||
torch.exp(self.scales[split_mask].repeat(samps, 1)) * centered_samples
|
||||
) # how these scales are rotated
|
||||
quats = self.quats[split_mask] / self.quats[split_mask].norm(dim=-1, keepdim=True) # normalize them first
|
||||
rots = quat_to_rotmat(quats.repeat(samps, 1)) # how these scales are rotated
|
||||
rotated_samples = torch.bmm(rots, scaled_samples[..., None]).squeeze()
|
||||
new_means = rotated_samples + self.means[split_mask].repeat(samps, 1)
|
||||
# step 2, sample new colors
|
||||
new_features_dc = self.features_dc[split_mask].repeat(samps, 1)
|
||||
new_features_rest = self.features_rest[split_mask].repeat(samps, 1, 1)
|
||||
# step 3, sample new opacities
|
||||
new_opacities = self.opacities[split_mask].repeat(samps, 1)
|
||||
# step 4, sample new scales
|
||||
size_fac = 1.6
|
||||
new_scales = torch.log(torch.exp(self.scales[split_mask]) / size_fac).repeat(samps, 1)
|
||||
self.scales[split_mask] = torch.log(torch.exp(self.scales[split_mask]) / size_fac)
|
||||
# step 5, sample new quats
|
||||
new_quats = self.quats[split_mask].repeat(samps, 1)
|
||||
out = {
|
||||
"means": new_means,
|
||||
"features_dc": new_features_dc,
|
||||
"features_rest": new_features_rest,
|
||||
"opacities": new_opacities,
|
||||
"scales": new_scales,
|
||||
"quats": new_quats,
|
||||
}
|
||||
for name, param in self.gauss_params.items():
|
||||
if name not in out:
|
||||
out[name] = param[split_mask].repeat(samps, 1)
|
||||
return out
|
||||
|
||||
def dup_gaussians(self, dup_mask):
|
||||
"""
|
||||
This function duplicates gaussians that are too small
|
||||
"""
|
||||
n_dups = dup_mask.sum().item()
|
||||
CONSOLE.log(f"Duplicating {dup_mask.sum().item()/self.num_points} gaussians: {n_dups}/{self.num_points}")
|
||||
new_dups = {}
|
||||
for name, param in self.gauss_params.items():
|
||||
new_dups[name] = param[dup_mask]
|
||||
return new_dups
|
||||
|
||||
def get_training_callbacks(
|
||||
self, training_callback_attributes: TrainingCallbackAttributes
|
||||
@ -777,7 +678,6 @@ class SplatfactoModel(Model):
|
||||
absgrad=True,
|
||||
rasterize_mode=self.config.rasterize_mode,
|
||||
# set some threshold to disregrad small gaussians for faster rendering.
|
||||
# radius_clip=3.0,
|
||||
)
|
||||
if self.training and info["means2d"].requires_grad:
|
||||
info["means2d"].retain_grad()
|
||||
@ -890,6 +790,8 @@ class SplatfactoModel(Model):
|
||||
if self.training:
|
||||
# Add loss from camera optimizer
|
||||
self.camera_optimizer.get_loss_dict(loss_dict)
|
||||
loss_dict['opacity_reg'] = .01 * torch.abs(torch.sigmoid(self.gauss_params["opacities"])).mean()
|
||||
loss_dict['scale_reg_mcmc'] = .01 * torch.abs(torch.exp(self.gauss_params["scales"])).mean()
|
||||
|
||||
return loss_dict
|
||||
|
||||
@ -939,4 +841,4 @@ class SplatfactoModel(Model):
|
||||
|
||||
images_dict = {"img": combined_rgb}
|
||||
|
||||
return metrics_dict, images_dict
|
||||
return metrics_dict, images_dict
|
||||
Loading…
Reference in New Issue
Block a user