add lite big models

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):

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)

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",
 )

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'