From bfdda8ca6ac5513eeecfd7545ffbce770df8ef80 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Mon, 1 Jul 2024 01:34:27 -0700 Subject: [PATCH] port gsplat impl over to splatfacto --- nerfstudio/models/splatfacto.py | 334 +++++++++++--------------------- 1 file changed, 118 insertions(+), 216 deletions(-) diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index fc9d85474..d0662d71d 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -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 \ No newline at end of file