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

Warmups

Browse Source
This commit is contained in:
Matt Tancik 2023-02-21 17:12:23 -08:00
parent 77d15f9b74
commit 207959da16
4 changed files with 27 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
nerfstudio/configs/method_configs.py
View File

@ -349,21 +349,22 @@ method_configs["dreamfusion"] = TrainerConfig(
dataparser=NerfstudioDataParserConfig(),
train_num_rays_per_batch=4096,
eval_num_rays_per_batch=4096,
horizontal_rotation_warmup=4000,
),
model=DreamFusionModelConfig(
eval_num_rays_per_chunk=1 << 15,
distortion_loss_mult=0.02,
interlevel_loss_mult=1.0,
interlevel_loss_mult=100.0,
orientation_loss_mult=0.1,
max_res=256,
sphere_collider=True,
initialize_density=False,
random_background=True,
proposal_warmup=1000,
proposal_warmup=500,
proposal_update_every=2,
proposal_weights_anneal_max_num_iters=100,
start_lambertian_training=1000,
start_normals_training=500,
start_lambertian_training=500,
start_normals_training=1000,
opacity_loss_mult=0.001,
positional_prompting="discrete",
guidance_scale=100,
@ -375,8 +376,8 @@ method_configs["dreamfusion"] = TrainerConfig(
"scheduler": None,
},
"fields": {
"optimizer": AdamOptimizerConfig(lr=5e-3, eps=1e-15),
"scheduler": None,
"optimizer": AdamOptimizerConfig(lr=1e-3, eps=1e-15),
"scheduler": ExponentialDecaySchedulerConfig(warmup_steps=2000, lr_final=1e-6, max_steps=20000),
},
},
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),

11
nerfstudio/data/datamanagers/dreamfusion_datamanager.py
View File

@ -66,7 +66,7 @@ def random_train_pose(
radius_mean: float = 1.0,
radius_std: float = 0.1,
central_rotation_range: Tuple[float, float] = (0, 360),
vertical_rotation_range: Tuple[float, float] = (-80, 20),
vertical_rotation_range: Tuple[float, float] = (-90, 0),
focal_range: Tuple[float, float] = (0.75, 1.35),
jitter_std: float = 0.01,
center: Tuple[float, float, float] = (0, 0, 0),
@ -169,11 +169,14 @@ class DreamFusionDataManagerConfig(VanillaDataManagerConfig):
"""Std of radius of camera orbit"""
focal_range: Tuple[float, float] = (0.6, 1.2)
"""Range of focal length"""
vertical_rotation_range: Tuple[float, float] = (-80, 30)
vertical_rotation_range: Tuple[float, float] = (-90, 0)
"""Range of vertical rotation"""
jitter_std: float = 0.05
"""Std of camera direction jitter, so we don't just point the cameras towards the center every time"""
center: Tuple[float, float, float] = (0, 0, 0)
"""Center coordinate of the camera sphere"""
horizontal_rotation_warmup: int = 0
"""How many steps until the full horizontal rotation range is used"""
class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-method
@ -244,7 +247,8 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me
# ).flatten()
# return ray_bundle, {"initialization": False}
# TODO below
horizontal_range = min((step / self.config.horizontal_rotation_warmup), 1) * 180
cameras, vertical_rotation, central_rotation = random_train_pose(
self.config.train_images_per_batch,
self.config.train_resolution,
@ -255,6 +259,7 @@ class DreamFusionDataManager(VanillaDataManager): # pylint: disable=abstract-me
vertical_rotation_range=self.config.vertical_rotation_range,
jitter_std=self.config.jitter_std,
center=self.config.center,
central_rotation_range=(-horizontal_range, horizontal_range),
)
ray_bundle = cameras.generate_rays(torch.tensor(list(range(self.config.train_images_per_batch)))).flatten()

2
nerfstudio/generative/stable_diffusion_utils.py
View File

@ -46,7 +46,7 @@ class PositionalTextEmbeddings:
back_prompt: str,
front_prompt: str,
stable_diffusion: StableDiffusion,
positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete",
positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete"
):
self.positional_prompting = positional_prompting
self.sd_device = stable_diffusion.device

21
nerfstudio/models/dreamfusion.py
View File

@ -77,8 +77,10 @@ class DreamFusionModelConfig(ModelConfig):
"""Randomizes light source per output."""
initialize_density: bool = True
"""Initialize density in center of scene."""
init_density_strength: float = 10
"""Initial strength of center density"""
taper_range: Tuple[int, int] = (0, 1000)
"""Range of step values for the density tapering"""
taper_strength: Tuple[float, float] = (1.0, 0.0)
"""Strength schedule of center density"""
sphere_collider: bool = True
"""Use spherical collider instead of box"""
random_background: bool = True
@ -126,7 +128,7 @@ class DreamFusionModelConfig(ModelConfig):
"""start training with lambertian shading after this many iterations"""
opacity_penalty: bool = True
"""enables penalty to encourage sparse weights (penalizing for uniform density along ray)"""
opacity_loss_mult: float = 1
opacity_loss_mult: float = 1e-3
"""scale for opacity penalty"""
max_res: int = 256
"""Maximum resolution of the density field."""
@ -135,7 +137,7 @@ class DreamFusionModelConfig(ModelConfig):
"""enables location based prompting"""
interpolated_prompting: bool = False
"""enables interpolated location prompting"""
positional_prompting: Literal["discrete", "interpolated", "base"] = "discrete"
positional_prompting: Literal["discrete", "interpolated", "off"] = "discrete"
""" how to incorporate position into prompt"""
top_prompt: str = ", overhead view"
"""appended to prompt for overhead view"""
@ -237,7 +239,7 @@ class DreamFusionModel(Model):
single_jitter=self.config.use_single_jitter,
update_sched=update_schedule,
initial_sampler=UniformSampler(single_jitter=self.config.use_single_jitter),
pdf_histogram_padding=0.001,
pdf_histogram_padding=0.005,
)
# renderers
@ -267,7 +269,7 @@ class DreamFusionModel(Model):
def taper_density(
self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument
):
self.density_strength -= 0.25
self.density_strength = np.interp(step, self.config.taper_range, self.config.taper_strength)
def start_training_normals(
self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument
@ -303,6 +305,7 @@ class DreamFusionModel(Model):
where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
iters=(150, 300, 450, 600),
func=taper_density,
update_every_num_iters=1,
args=[self, training_callback_attributes],
),
TrainingCallback(
@ -352,11 +355,9 @@ class DreamFusionModel(Model):
if self.initialize_density:
pos = ray_samples.frustums.get_positions()
density_blob = (
self.config.init_density_strength
* self.density_strength
* torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None]
self.density_strength * self.density_strength * torch.exp(-torch.norm(pos, dim=-1) * 10)[..., None]
)
density += density_blob
density = density + density_blob
weights = ray_samples.get_weights(density)
weights_list.append(weights)