tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Fix opacity loss, use unform sampling first round of prop

Browse Source
This commit is contained in:
Matt Tancik 2023-02-15 13:43:49 -08:00
parent 06ec4bd982
commit 4db2410b40
3 changed files with 30 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
nerfstudio/configs/method_configs.py
View File

@ -341,7 +341,7 @@ method_configs["dreamfusion"] = TrainerConfig(
method_name="dreamfusion",
steps_per_eval_batch=50,
steps_per_eval_image=50,
steps_per_save=2000,
steps_per_save=200,
max_num_iterations=30000,
mixed_precision=True,
pipeline=DreamfusionPipelineConfig(
@ -353,7 +353,7 @@ method_configs["dreamfusion"] = TrainerConfig(
model=DreamFusionModelConfig(
eval_num_rays_per_chunk=1 << 15,
distortion_loss_mult=10.0,
orientation_loss_mult=0.0001,
orientation_loss_mult=0.1,
max_res=256,
sphere_collider=True,
initialize_density=False,
@ -364,6 +364,7 @@ method_configs["dreamfusion"] = TrainerConfig(
proposal_weights_anneal_max_num_iters=100,
start_lambertian_training=1000,
start_normals_training=500,
opacity_loss_mult=0.001,
),
interpolated_prompting=False,
guidance_scale=100,

17
nerfstudio/model_components/ray_samplers.py
View File

@ -509,7 +509,16 @@ class VolumetricSampler(Sampler):
class ProposalNetworkSampler(Sampler):
"""Sampler that uses a proposal network to generate samples."""
"""Sampler that uses a proposal network to generate samples.
Args:
num_proposal_samples_per_ray: Number of samples to generate per ray for each proposal step.
num_nerf_samples_per_ray: Number of samples to generate per ray for the NERF model.
num_proposal_network_iterations: Number of proposal network iterations to run.
single_jitter: Use a same random jitter for all samples along a ray.
update_sched: A function that takes the iteration number of steps between updates.
initial_sampler: Sampler to use for the first iteration. Uses UniformLinDispPiecewise if not set.
"""
def __init__(
self,
@ -518,6 +527,7 @@ class ProposalNetworkSampler(Sampler):
num_proposal_network_iterations: int = 2,
single_jitter: bool = False,
update_sched: Callable = lambda x: 1,
initial_sampler: Optional[Sampler] = None,
) -> None:
super().__init__()
self.num_proposal_samples_per_ray = num_proposal_samples_per_ray
@ -528,7 +538,10 @@ class ProposalNetworkSampler(Sampler):
raise ValueError("num_proposal_network_iterations must be >= 1")
# samplers
self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
if initial_sampler is None:
self.initial_sampler = UniformLinDispPiecewiseSampler(single_jitter=single_jitter)
else:
self.initial_sampler = initial_sampler
self.pdf_sampler = PDFSampler(include_original=False, single_jitter=single_jitter)
self._anneal = 1.0

29
nerfstudio/models/dreamfusion.py
View File

@ -41,7 +41,10 @@ from nerfstudio.model_components.losses import (
orientation_loss,
pred_normal_loss,
)
from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler
from nerfstudio.model_components.ray_samplers import (
ProposalNetworkSampler,
UniformSampler,
)
from nerfstudio.model_components.renderers import (
AccumulationRenderer,
DepthRenderer,
@ -116,12 +119,8 @@ class DreamFusionModelConfig(ModelConfig):
"""Start training normals after this many iterations"""
start_lambertian_training: int = 1000
"""start training with lambertian shading after this many iterations"""
alphas_penalty: bool = True
"""enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
alphas_loss_mult: float = 1
"""scale for alphas penalty"""
opacity_penalty: bool = True
"""enables penalty to encourage transparent scenes, as in "dreamfields" paper"""
"""enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
opacity_loss_mult: float = 1
"""scale for opacity penalty"""
max_res: int = 256
@ -183,10 +182,9 @@ class DreamFusionModel(Model):
num_proposal_network_iterations=self.config.num_proposal_iterations,
single_jitter=self.config.use_single_jitter,
update_sched=update_schedule,
initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter),
)
# self.sampler_uniform = UniformSampler(num_samples=self.num_samples, single_jitter=True)
# renderers
self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
self.renderer_accumulation = AccumulationRenderer()
@ -373,7 +371,7 @@ class DreamFusionModel(Model):
if self.training or True:
outputs["rendered_orientation_loss"] = orientation_loss(
weights, field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
)
outputs["rendered_pred_normal_loss"] = pred_normal_loss(
@ -383,9 +381,9 @@ class DreamFusionModel(Model):
)
assert weights.shape[-1] == 1
if self.config.alphas_penalty:
outputs["alphas_loss"] = (
torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.alphas_loss_mult
if self.config.opacity_penalty:
outputs["opacity_loss"] = (
torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) * self.config.opacity_loss_mult
)
return outputs
@ -408,12 +406,9 @@ class DreamFusionModel(Model):
loss_dict["orientation_loss"] = 0
loss_dict["pred_normal_loss"] = 0
if self.config.alphas_penalty:
loss_dict["alphas_loss"] = self.config.alphas_loss_mult * outputs["alphas_loss"].mean()
if self.config.opacity_penalty:
opacity_loss = -torch.minimum((1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance))
loss_dict["opacity_loss"] = self.config.opacity_loss_mult * opacity_loss
loss_dict["opacity_loss"] = self.config.opacity_loss_mult * outputs["opacity_loss"].mean()
if self.training:
loss_dict["distortion_loss"] = self.config.distortion_loss_mult * distortion_loss(
outputs["weights_list"], outputs["ray_samples_list"]