diff --git a/projects-lerf/lerf/encoders/openclip_encoder.py b/projects-lerf/lerf/encoders/openclip_encoder.py index 559359be5..1a6dbb29e 100644 --- a/projects-lerf/lerf/encoders/openclip_encoder.py +++ b/projects-lerf/lerf/encoders/openclip_encoder.py @@ -11,6 +11,7 @@ except ImportError: from lerf.encoders.image_encoder import BaseImageEncoder, BaseImageEncoderConfig + @dataclass class OpenCLIPNetworkConfig(BaseImageEncoderConfig): _target: Type = field(default_factory=lambda: OpenCLIPNetwork) @@ -20,9 +21,6 @@ class OpenCLIPNetworkConfig(BaseImageEncoderConfig): clip_model_type: str = "ViT-B-16" clip_model_pretrained: str = "laion2b_s34b_b88k" clip_n_dims: int = 512 - # clip_model_type: str = "ViT-L-14" - # clip_model_pretrained: str = "laion2b_s32b_b82k" - # clip_n_dims: int = 768 negatives: Tuple[str] = ("object", "things", "stuff", "texture") class OpenCLIPNetwork(BaseImageEncoder): diff --git a/projects-lerf/lerf/lerf.py b/projects-lerf/lerf/lerf.py index 6932f360b..8e53cfbd6 100644 --- a/projects-lerf/lerf/lerf.py +++ b/projects-lerf/lerf/lerf.py @@ -28,7 +28,7 @@ class LERFModelConfig(NerfactoModelConfig): max_scale: float = 1.5 """maximum scale used to compute relevancy with""" specify_scale: bool = False - num_lerf_samples: int = 12 + num_lerf_samples: int = 24 hashgrid_layers: Tuple[int] = (12, 12) hashgrid_resolutions: Tuple[Tuple[int]] = ((16, 128), (128, 512)) hashgrid_sizes: Tuple[int] = (19, 19) diff --git a/projects-lerf/lerf/lerf_config.py b/projects-lerf/lerf/lerf_config.py index 73d203259..3b4997b0e 100644 --- a/projects-lerf/lerf/lerf_config.py +++ b/projects-lerf/lerf/lerf_config.py @@ -41,6 +41,7 @@ lerf_method = MethodSpecification( hashgrid_sizes=(19, 19), hashgrid_layers=(12, 12), hashgrid_resolutions=((16, 128), (128, 512)), + num_lerf_samples=24, ), network=OpenCLIPNetworkConfig( clip_model_type="ViT-B-16", clip_model_pretrained="laion2b_s34b_b88k", clip_n_dims=512 @@ -63,5 +64,98 @@ lerf_method = MethodSpecification( viewer=ViewerConfig(num_rays_per_chunk=1 << 15), vis="viewer", ), - description="TODO", + description="Base config for LERF", +) +lerf_method_big = MethodSpecification( + config=TrainerConfig( + method_name="lerf-big", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=LERFPipelineConfig( + datamanager=LERFDataManagerConfig( + dataparser=NerfstudioDataParserConfig(train_split_fraction=0.99), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=LERFModelConfig( + eval_num_rays_per_chunk=1 << 16, + hashgrid_sizes=(19, 19), + hashgrid_layers=(14, 14), + hashgrid_resolutions=((16, 128), (128, 512)), + num_lerf_samples=32, + ), + network=OpenCLIPNetworkConfig( + clip_model_type="ViT-L-14", clip_model_pretrained="laion2b_s32b_b82k", clip_n_dims=768 + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=30000), + }, + "lerf": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=7000), + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", + ), + description="A larger version of LERF with a higher memory footprint, bigger CLIP model, and more hashgrid capacity", +) + +lerf_method_lite = MethodSpecification( + config=TrainerConfig( + method_name="lerf-lite", + steps_per_eval_batch=500, + steps_per_save=2000, + max_num_iterations=30000, + mixed_precision=True, + pipeline=LERFPipelineConfig( + datamanager=LERFDataManagerConfig( + dataparser=NerfstudioDataParserConfig(train_split_fraction=0.99), + train_num_rays_per_batch=4096, + eval_num_rays_per_batch=4096, + camera_optimizer=CameraOptimizerConfig( + mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) + ), + ), + model=LERFModelConfig( + eval_num_rays_per_chunk=1 << 16, + hashgrid_sizes=(19,), + hashgrid_layers=(16,), + hashgrid_resolutions=((16, 512),), + num_lerf_samples=12, + ), + network=OpenCLIPNetworkConfig( + clip_model_type="ViT-B-16", clip_model_pretrained="laion2b_s34b_b88k", clip_n_dims=512 + ), + ), + optimizers={ + "proposal_networks": { + "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": None, + }, + "fields": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=30000), + }, + "lerf": { + "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-3, max_steps=7000), + }, + }, + viewer=ViewerConfig(num_rays_per_chunk=1 << 15), + vis="viewer", + ), + description="A lightweight version of LERF designed to work on smaller GPUs", ) diff --git a/projects-lerf/pyproject.toml b/projects-lerf/pyproject.toml index aff2088c1..f2d251593 100644 --- a/projects-lerf/pyproject.toml +++ b/projects-lerf/pyproject.toml @@ -4,3 +4,5 @@ version = "0.1.0" [project.entry-points.'nerfstudio.method_configs'] lerf = 'lerf.lerf_config:lerf_method' +lerf_big = 'lerf.lerf_config:lerf_method_big' +lerf_lite = 'lerf.lerf_config:lerf_method_lite' \ No newline at end of file