[WIP] Fix model loading collate_fn non-static (#1990)

* Fix model loading collate_fn non-static

* Fix serialization

* Add tests to reproduce the error

* Comment the change to let the tests fail

* Finish changes

nerfstudio/data/datamanagers/base_datamanager.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 
 import functools
 from abc import abstractmethod
+from collections import defaultdict
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import (
@@ -33,6 +34,7 @@ from typing import (
     Tuple,
     Type,
     Union,
+    cast,
 )
 
 import torch
@@ -64,7 +66,6 @@ from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttrib
 from nerfstudio.model_components.ray_generators import RayGenerator
 from nerfstudio.utils.misc import IterableWrapper
 from nerfstudio.utils.rich_utils import CONSOLE
-from collections import defaultdict
 
 
 def variable_res_collate(batch: List[Dict]) -> Dict:
@@ -331,7 +332,7 @@ class VanillaDataManagerConfig(DataManagerConfig):
     camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig()
     """Specifies the camera pose optimizer used during training. Helpful if poses are noisy, such as for data from
     Record3D."""
-    collate_fn: Callable[[Any], Any] = nerfstudio_collate
+    collate_fn: Callable[[Any], Any] = cast(Any, staticmethod(nerfstudio_collate))
     """Specifies the collate function to use for the train and eval dataloaders."""
     camera_res_scale_factor: float = 1.0
     """The scale factor for scaling spatial data such as images, mask, semantics

tests/data/configs/test_config1.yml

@@ -0,0 +1,194 @@
+!!python/object:nerfstudio.engine.trainer.TrainerConfig
+_target: !!python/name:nerfstudio.engine.trainer.Trainer ''
+data: null
+experiment_name: unnamed
+load_checkpoint: !!python/object/apply:pathlib.PosixPath
+- outputs
+- unnamed
+- nerfacto
+- 2023-05-24_204526
+- config.yml
+load_config: null
+load_dir: null
+load_step: null
+log_gradients: false
+logging: !!python/object:nerfstudio.configs.base_config.LoggingConfig
+  local_writer: !!python/object:nerfstudio.configs.base_config.LocalWriterConfig
+    _target: !!python/name:nerfstudio.utils.writer.LocalWriter ''
+    enable: true
+    max_log_size: 10
+    stats_to_track: !!python/tuple
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Train Iter (time)
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Train Rays / Sec
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Test PSNR
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Vis Rays / Sec
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Test Rays / Sec
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - ETA (time)
+  max_buffer_size: 20
+  profiler: basic
+  relative_log_dir: !!python/object/apply:pathlib.PosixPath []
+  steps_per_log: 10
+machine: !!python/object:nerfstudio.configs.base_config.MachineConfig
+  dist_url: auto
+  machine_rank: 0
+  num_gpus: 1
+  num_machines: 1
+  seed: 42
+max_num_iterations: 30000
+method_name: nerfacto
+mixed_precision: true
+optimizers:
+  fields:
+    optimizer: !!python/object:nerfstudio.engine.optimizers.AdamOptimizerConfig
+      _target: &id001 !!python/name:torch.optim.adam.Adam ''
+      eps: 1.0e-15
+      lr: 0.01
+      max_norm: null
+      weight_decay: 0
+    scheduler: !!python/object:nerfstudio.engine.schedulers.ExponentialDecaySchedulerConfig
+      _target: &id002 !!python/name:nerfstudio.engine.schedulers.ExponentialDecayScheduler ''
+      lr_final: 0.0001
+      lr_pre_warmup: 1.0e-08
+      max_steps: 200000
+      ramp: cosine
+      warmup_steps: 0
+  proposal_networks:
+    optimizer: !!python/object:nerfstudio.engine.optimizers.AdamOptimizerConfig
+      _target: *id001
+      eps: 1.0e-15
+      lr: 0.01
+      max_norm: null
+      weight_decay: 0
+    scheduler: !!python/object:nerfstudio.engine.schedulers.ExponentialDecaySchedulerConfig
+      _target: *id002
+      lr_final: 0.0001
+      lr_pre_warmup: 1.0e-08
+      max_steps: 200000
+      ramp: cosine
+      warmup_steps: 0
+output_dir: !!python/object/apply:pathlib.PosixPath
+- outputs
+pipeline: !!python/object:nerfstudio.pipelines.base_pipeline.VanillaPipelineConfig
+  _target: !!python/name:nerfstudio.pipelines.base_pipeline.VanillaPipeline ''
+  datamanager: !!python/object:nerfstudio.data.datamanagers.base_datamanager.VanillaDataManagerConfig
+    _target: !!python/name:nerfstudio.data.datamanagers.base_datamanager.VanillaDataManager ''
+    camera_optimizer: !!python/object:nerfstudio.cameras.camera_optimizers.CameraOptimizerConfig
+      _target: !!python/name:nerfstudio.cameras.camera_optimizers.CameraOptimizer ''
+      mode: 'off'
+      optimizer: !!python/object:nerfstudio.engine.optimizers.AdamOptimizerConfig
+        _target: *id001
+        eps: 1.0e-08
+        lr: 0.0006
+        max_norm: null
+        weight_decay: 0.01
+      orientation_noise_std: 0.0
+      param_group: camera_opt
+      position_noise_std: 0.0
+      scheduler: !!python/object:nerfstudio.engine.schedulers.ExponentialDecaySchedulerConfig
+        _target: *id002
+        lr_final: 6.0e-06
+        lr_pre_warmup: 1.0e-08
+        max_steps: 200000
+        ramp: cosine
+        warmup_steps: 0
+    camera_res_scale_factor: 1.0
+    collate_fn: !!python/name:nerfstudio.data.utils.nerfstudio_collate.nerfstudio_collate ''
+    data: null
+    dataparser: !!python/object:nerfstudio.data.dataparsers.minimal_dataparser.MinimalDataParserConfig
+      _target: !!python/name:nerfstudio.data.dataparsers.minimal_dataparser.MinimalDataParser ''
+      data: !!python/object/apply:pathlib.PosixPath
+      - /
+      - home
+      - kulhajon
+      - projects
+      - nerf-playground
+      - data
+      - varun
+      - cheezit-scaled
+    eval_image_indices: !!python/tuple
+    - 0
+    eval_num_images_to_sample_from: -1
+    eval_num_rays_per_batch: 4096
+    eval_num_times_to_repeat_images: -1
+    patch_size: 1
+    train_num_images_to_sample_from: -1
+    train_num_rays_per_batch: 4096
+    train_num_times_to_repeat_images: -1
+  model: !!python/object:nerfstudio.models.nerfacto.NerfactoModelConfig
+    _target: !!python/name:nerfstudio.models.nerfacto.NerfactoModel ''
+    background_color: last_sample
+    collider_params:
+      far_plane: 6.0
+      near_plane: 2.0
+    disable_scene_contraction: false
+    distortion_loss_mult: 0.002
+    enable_collider: true
+    eval_num_rays_per_chunk: 32768
+    far_plane: 1000.0
+    hidden_dim: 64
+    hidden_dim_color: 64
+    hidden_dim_transient: 64
+    interlevel_loss_mult: 1.0
+    log2_hashmap_size: 19
+    loss_coefficients:
+      rgb_loss_coarse: 1.0
+      rgb_loss_fine: 1.0
+    max_res: 2048
+    near_plane: 0.05
+    num_levels: 16
+    num_nerf_samples_per_ray: 48
+    num_proposal_iterations: 2
+    num_proposal_samples_per_ray: !!python/tuple
+    - 256
+    - 96
+    orientation_loss_mult: 0.0001
+    pred_normal_loss_mult: 0.001
+    predict_normals: false
+    proposal_initial_sampler: piecewise
+    proposal_net_args_list:
+    - hidden_dim: 16
+      log2_hashmap_size: 17
+      max_res: 128
+      num_levels: 5
+      use_linear: false
+    - hidden_dim: 16
+      log2_hashmap_size: 17
+      max_res: 256
+      num_levels: 5
+      use_linear: false
+    proposal_update_every: 5
+    proposal_warmup: 5000
+    proposal_weights_anneal_max_num_iters: 1000
+    proposal_weights_anneal_slope: 10.0
+    use_average_appearance_embedding: true
+    use_gradient_scaling: false
+    use_proposal_weight_anneal: true
+    use_same_proposal_network: false
+    use_single_jitter: true
+project_name: nerfstudio-project
+relative_model_dir: !!python/object/apply:pathlib.PosixPath
+- nerfstudio_models
+save_only_latest_checkpoint: true
+steps_per_eval_all_images: 25000
+steps_per_eval_batch: 500
+steps_per_eval_image: 500
+steps_per_save: 2000
+timestamp: 2023-05-24_205104
+use_grad_scaler: false
+viewer: !!python/object:nerfstudio.configs.base_config.ViewerConfig
+  image_format: jpeg
+  jpeg_quality: 90
+  max_num_display_images: 512
+  num_rays_per_chunk: 32768
+  quit_on_train_completion: false
+  relative_log_filename: viewer_log_filename.txt
+  websocket_host: 0.0.0.0
+  websocket_port: null
+  websocket_port_default: 7007
+vis: viewer+wandb

tests/data/configs/test_config2.yml

@@ -0,0 +1,188 @@
+!!python/object:nerfstudio.engine.trainer.TrainerConfig
+_target: !!python/name:nerfstudio.engine.trainer.Trainer ''
+data: null
+experiment_name: unnamed
+load_checkpoint: null
+load_config: null
+load_dir: null
+load_step: null
+log_gradients: false
+logging: !!python/object:nerfstudio.configs.base_config.LoggingConfig
+  local_writer: !!python/object:nerfstudio.configs.base_config.LocalWriterConfig
+    _target: !!python/name:nerfstudio.utils.writer.LocalWriter ''
+    enable: true
+    max_log_size: 10
+    stats_to_track: !!python/tuple
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Train Iter (time)
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Train Rays / Sec
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Test PSNR
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Vis Rays / Sec
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - Test Rays / Sec
+    - !!python/object/apply:nerfstudio.utils.writer.EventName
+      - ETA (time)
+  max_buffer_size: 20
+  profiler: basic
+  relative_log_dir: !!python/object/apply:pathlib.PosixPath []
+  steps_per_log: 10
+machine: !!python/object:nerfstudio.configs.base_config.MachineConfig
+  dist_url: auto
+  machine_rank: 0
+  num_gpus: 1
+  num_machines: 1
+  seed: 42
+max_num_iterations: 30000
+method_name: nerfacto
+mixed_precision: true
+optimizers:
+  fields:
+    optimizer: !!python/object:nerfstudio.engine.optimizers.AdamOptimizerConfig
+      _target: &id001 !!python/name:torch.optim.adam.Adam ''
+      eps: 1.0e-15
+      lr: 0.01
+      max_norm: null
+      weight_decay: 0
+    scheduler: !!python/object:nerfstudio.engine.schedulers.ExponentialDecaySchedulerConfig
+      _target: &id002 !!python/name:nerfstudio.engine.schedulers.ExponentialDecayScheduler ''
+      lr_final: 0.0001
+      lr_pre_warmup: 1.0e-08
+      max_steps: 200000
+      ramp: cosine
+      warmup_steps: 0
+  proposal_networks:
+    optimizer: !!python/object:nerfstudio.engine.optimizers.AdamOptimizerConfig
+      _target: *id001
+      eps: 1.0e-15
+      lr: 0.01
+      max_norm: null
+      weight_decay: 0
+    scheduler: !!python/object:nerfstudio.engine.schedulers.ExponentialDecaySchedulerConfig
+      _target: *id002
+      lr_final: 0.0001
+      lr_pre_warmup: 1.0e-08
+      max_steps: 200000
+      ramp: cosine
+      warmup_steps: 0
+output_dir: !!python/object/apply:pathlib.PosixPath
+- outputs
+pipeline: !!python/object:nerfstudio.pipelines.base_pipeline.VanillaPipelineConfig
+  _target: !!python/name:nerfstudio.pipelines.base_pipeline.VanillaPipeline ''
+  datamanager: !!python/object:nerfstudio.data.datamanagers.base_datamanager.VanillaDataManagerConfig
+    _target: !!python/name:nerfstudio.data.datamanagers.base_datamanager.VanillaDataManager ''
+    camera_optimizer: !!python/object:nerfstudio.cameras.camera_optimizers.CameraOptimizerConfig
+      _target: !!python/name:nerfstudio.cameras.camera_optimizers.CameraOptimizer ''
+      mode: 'off'
+      optimizer: !!python/object:nerfstudio.engine.optimizers.AdamOptimizerConfig
+        _target: *id001
+        eps: 1.0e-08
+        lr: 0.0006
+        max_norm: null
+        weight_decay: 0.01
+      orientation_noise_std: 0.0
+      param_group: camera_opt
+      position_noise_std: 0.0
+      scheduler: !!python/object:nerfstudio.engine.schedulers.ExponentialDecaySchedulerConfig
+        _target: *id002
+        lr_final: 6.0e-06
+        lr_pre_warmup: 1.0e-08
+        max_steps: 200000
+        ramp: cosine
+        warmup_steps: 0
+    camera_res_scale_factor: 1.0
+    data: null
+    dataparser: !!python/object:nerfstudio.data.dataparsers.minimal_dataparser.MinimalDataParserConfig
+      _target: !!python/name:nerfstudio.data.dataparsers.minimal_dataparser.MinimalDataParser ''
+      data: !!python/object/apply:pathlib.PosixPath
+      - /
+      - home
+      - kulhajon
+      - projects
+      - nerf-playground
+      - data
+      - varun
+      - cheezit-scaled
+    eval_image_indices: !!python/tuple
+    - 0
+    eval_num_images_to_sample_from: -1
+    eval_num_rays_per_batch: 4096
+    eval_num_times_to_repeat_images: -1
+    patch_size: 1
+    train_num_images_to_sample_from: -1
+    train_num_rays_per_batch: 4096
+    train_num_times_to_repeat_images: -1
+  model: !!python/object:nerfstudio.models.nerfacto.NerfactoModelConfig
+    _target: !!python/name:nerfstudio.models.nerfacto.NerfactoModel ''
+    background_color: last_sample
+    collider_params:
+      far_plane: 6.0
+      near_plane: 2.0
+    disable_scene_contraction: false
+    distortion_loss_mult: 0.002
+    enable_collider: true
+    eval_num_rays_per_chunk: 32768
+    far_plane: 1000.0
+    hidden_dim: 64
+    hidden_dim_color: 64
+    hidden_dim_transient: 64
+    interlevel_loss_mult: 1.0
+    log2_hashmap_size: 19
+    loss_coefficients:
+      rgb_loss_coarse: 1.0
+      rgb_loss_fine: 1.0
+    max_res: 2048
+    near_plane: 0.05
+    num_levels: 16
+    num_nerf_samples_per_ray: 48
+    num_proposal_iterations: 2
+    num_proposal_samples_per_ray: !!python/tuple
+    - 256
+    - 96
+    orientation_loss_mult: 0.0001
+    pred_normal_loss_mult: 0.001
+    predict_normals: false
+    proposal_initial_sampler: piecewise
+    proposal_net_args_list:
+    - hidden_dim: 16
+      log2_hashmap_size: 17
+      max_res: 128
+      num_levels: 5
+      use_linear: false
+    - hidden_dim: 16
+      log2_hashmap_size: 17
+      max_res: 256
+      num_levels: 5
+      use_linear: false
+    proposal_update_every: 5
+    proposal_warmup: 5000
+    proposal_weights_anneal_max_num_iters: 1000
+    proposal_weights_anneal_slope: 10.0
+    use_average_appearance_embedding: true
+    use_gradient_scaling: false
+    use_proposal_weight_anneal: true
+    use_same_proposal_network: false
+    use_single_jitter: true
+project_name: nerfstudio-project
+relative_model_dir: !!python/object/apply:pathlib.PosixPath
+- nerfstudio_models
+save_only_latest_checkpoint: true
+steps_per_eval_all_images: 25000
+steps_per_eval_batch: 500
+steps_per_eval_image: 500
+steps_per_save: 2000
+timestamp: 2023-05-24_204526
+use_grad_scaler: false
+viewer: !!python/object:nerfstudio.configs.base_config.ViewerConfig
+  image_format: jpeg
+  jpeg_quality: 90
+  max_num_display_images: 512
+  num_rays_per_chunk: 32768
+  quit_on_train_completion: false
+  relative_log_filename: viewer_log_filename.txt
+  websocket_host: 0.0.0.0
+  websocket_port: null
+  websocket_port_default: 7007
+vis: viewer+wandb

tests/data/test_datamanager.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import Any
 
 import pytest
@@ -102,3 +103,21 @@ def test_data_manager_type_can_be_serialized(config):
         assert isinstance(obj, tmp2)
     finally:
         globals().pop("tmp2")
+
+
+def _dummy_function():
+    return True
+
+
+def test_deserialize_config1():
+    with open(Path(__file__).parent / "configs" / "test_config1.yml", "r") as f:
+        config_str = f.read()
+    obj = yaml.load(config_str, Loader=yaml.Loader)
+    obj.pipeline.datamanager.collate_fn([1, 2, 3])
+
+
+def test_deserialize_config2():
+    with open(Path(__file__).parent / "configs" / "test_config2.yml", "r") as f:
+        config_str = f.read()
+    obj = yaml.load(config_str, Loader=yaml.Loader)
+    obj.pipeline.datamanager.collate_fn([1, 2, 3])