diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index d2db011d3..a03afa5f9 100644 --- a/nerfstudio/model_components/losses.py +++ b/nerfstudio/model_components/losses.py @@ -120,6 +120,69 @@ def interlevel_loss(weights_list, ray_samples_list): return loss_interlevel +def blur_stepfun(x, y, r): + xr, xr_idx = torch.sort(torch.cat([x - r, x + r], dim=-1)) + y1 = ( + torch.cat([y, torch.zeros_like(y[..., :1])], dim=-1) - torch.cat([torch.zeros_like(y[..., :1]), y], dim=-1) + ) / (2 * r) + y2 = torch.cat([y1, -y1], dim=-1).take_along_dim(xr_idx[..., :-1], dim=-1) + yr = torch.cumsum((xr[..., 1:] - xr[..., :-1]) * torch.cumsum(y2, dim=-1), dim=-1).clamp_min(0) + yr = torch.cat([torch.zeros_like(yr[..., :1]), yr], dim=-1) + return xr, yr + + +def sorted_interp_quad(x, xp, fpdf, fcdf): + """interp in quadratic""" + + # Identify the location in `xp` that corresponds to each `x`. + # The final `True` index in `mask` is the start of the matching interval. + mask = x[..., None, :] >= xp[..., :, None] + + def find_interval(x): + # Grab the value where `mask` switches from True to False, and vice versa. + # This approach takes advantage of the fact that `x` is sorted. + x0 = torch.max(torch.where(mask, x[..., None], x[..., :1, None]), -2).values + x1 = torch.min(torch.where(~mask, x[..., None], x[..., -1:, None]), -2).values + return x0, x1 + + fpdf0, fpdf1 = find_interval(fpdf) + fcdf0, fcdf1 = find_interval(fcdf) + xp0, xp1 = find_interval(xp) + + offset = torch.clip(torch.nan_to_num((x - xp0) / (xp1 - xp0), 0), 0, 1) + ret = fcdf0 + (x - xp0) * (fpdf0 + fpdf1 * offset + fpdf0 * (1 - offset)) / 2 + return ret + + +def zipnerf_loss(weights_list, ray_samples_list): + # ground truth s and w (real nerf samples) + # This implementation matches ZipNeRF up to the scale. + # In the paper the loss is computed as the sum over the ray samples. + # Here we take the mean and the multiplier for this loss should be changed accordingly. + pulse_widths = [0.003, 0.03] + c = ray_samples_to_sdist(ray_samples_list[-1]).detach() + w = weights_list[-1][..., 0].detach() + + w_norm = w / (c[..., 1:] - c[..., :-1]) + loss = 0 + for i, (ray_samples, weights) in enumerate(zip(ray_samples_list[:-1], weights_list[:-1])): + cp = ray_samples_to_sdist(ray_samples) + wp = weights[..., 0] # (num_rays, num_samples) + c_, w_ = blur_stepfun(c, w_norm, pulse_widths[i]) + + # piecewise linear pdf to piecewise quadratic cdf + area = 0.5 * (w_[..., 1:] + w_[..., :-1]) * (c_[..., 1:] - c_[..., :-1]) + cdf = torch.cat([torch.zeros_like(area[..., :1]), torch.cumsum(area, dim=-1)], dim=-1) + + # query piecewise quadratic interpolation + cdf_interp = sorted_interp_quad(cp, c_, w_, cdf) + + # difference between adjacent interpolated values + w_s = torch.diff(cdf_interp, dim=-1) + loss += ((w_s - wp).clamp_min(0) ** 2 / (wp + 1e-5)).mean() + return loss + + # Verified def lossfun_distortion(t, w): """ diff --git a/nerfstudio/model_components/ray_samplers.py b/nerfstudio/model_components/ray_samplers.py index 7b7219f15..131db8308 100644 --- a/nerfstudio/model_components/ray_samplers.py +++ b/nerfstudio/model_components/ray_samplers.py @@ -262,7 +262,7 @@ class PDFSampler(Sampler): train_stratified: bool = True, single_jitter: bool = False, include_original: bool = True, - histogram_padding: float = 0.01, + histogram_padding: float = 0.00, ) -> None: super().__init__(num_samples=num_samples) self.train_stratified = train_stratified @@ -384,7 +384,6 @@ class VolumetricSampler(Sampler): occupancy_grid: OccGridEstimator, density_fn: Optional[Callable[[TensorType[..., 3]], TensorType[..., 1]]] = None, ) -> None: - super().__init__() assert occupancy_grid is not None self.density_fn = density_fn @@ -651,7 +650,6 @@ class NeuSSampler(Sampler): base_variance = self.base_variance while total_iters < self.num_upsample_steps: - with torch.no_grad(): new_sdf = sdf_fn(new_samples) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 53173640a..b3dcb9e76 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -45,6 +45,7 @@ from nerfstudio.model_components.losses import ( interlevel_loss, orientation_loss, pred_normal_loss, + zipnerf_loss, ) from nerfstudio.model_components.ray_samplers import ( ProposalNetworkSampler, @@ -85,9 +86,9 @@ class NerfactoModelConfig(ModelConfig): """Maximum resolution of the hashmap for the base mlp.""" log2_hashmap_size: int = 19 """Size of the hashmap for the base mlp""" - num_proposal_samples_per_ray: Tuple[int, ...] = (256, 96) + num_proposal_samples_per_ray: Tuple[int, ...] = (256, 128) """Number of samples per ray for each proposal network.""" - num_nerf_samples_per_ray: int = 48 + num_nerf_samples_per_ray: int = 64 """Number of samples per ray for the nerf network.""" proposal_update_every: int = 5 """Sample every n steps after the warmup""" @@ -99,8 +100,8 @@ class NerfactoModelConfig(ModelConfig): """Use the same proposal network. Otherwise use different ones.""" proposal_net_args_list: List[Dict] = field( default_factory=lambda: [ - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": False}, {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": False}, + {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False}, ] ) """Arguments for the proposal density fields.""" @@ -120,7 +121,7 @@ class NerfactoModelConfig(ModelConfig): """Whether to use average appearance embedding or zeros for inference.""" proposal_weights_anneal_slope: float = 10.0 """Slope of the annealing function for the proposal weights.""" - proposal_weights_anneal_max_num_iters: int = 1000 + proposal_weights_anneal_max_num_iters: int = 1 """Max num iterations for the annealing function.""" use_single_jitter: bool = True """Whether use single jitter or not for the proposal networks.""" @@ -128,6 +129,8 @@ class NerfactoModelConfig(ModelConfig): """Whether to predict normals or not.""" disable_scene_contraction: bool = False """Whether to disable scene contraction or not.""" + use_zipnerf_loss: bool = False + """Whether to use zipnerf loss or not.""" class NerfactoModel(Model): @@ -318,9 +321,14 @@ class NerfactoModel(Model): image = batch["image"].to(self.device) loss_dict["rgb_loss"] = self.rgb_loss(image, outputs["rgb"]) if self.training: - loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( - outputs["weights_list"], outputs["ray_samples_list"] - ) + if self.config.use_zipnerf_loss: + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * zipnerf_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) + else: + loss_dict["interlevel_loss"] = self.config.interlevel_loss_mult * interlevel_loss( + outputs["weights_list"], outputs["ray_samples_list"] + ) assert metrics_dict is not None and "distortion" in metrics_dict loss_dict["distortion_loss"] = self.config.distortion_loss_mult * metrics_dict["distortion"] if self.config.predict_normals: