From a0d5307dd788294a4d818e0a7d016438c5eda3a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= Date: Mon, 29 May 2023 09:56:41 +0200 Subject: [PATCH] [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 --- .../data/datamanagers/base_datamanager.py | 5 +- tests/data/configs/test_config1.yml | 194 ++++++++++++++++++ tests/data/configs/test_config2.yml | 188 +++++++++++++++++ tests/data/test_datamanager.py | 19 ++ 4 files changed, 404 insertions(+), 2 deletions(-) create mode 100644 tests/data/configs/test_config1.yml create mode 100644 tests/data/configs/test_config2.yml diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 707659d95..aeea43e3c 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/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 diff --git a/tests/data/configs/test_config1.yml b/tests/data/configs/test_config1.yml new file mode 100644 index 000000000..d58416491 --- /dev/null +++ b/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 diff --git a/tests/data/configs/test_config2.yml b/tests/data/configs/test_config2.yml new file mode 100644 index 000000000..92aadf1fd --- /dev/null +++ b/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 diff --git a/tests/data/test_datamanager.py b/tests/data/test_datamanager.py index 568e6de8f..e1a2a24c0 100644 --- a/tests/data/test_datamanager.py +++ b/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])