diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 899833f83..03d048e3e 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/nerfstudio/cameras/camera_optimizers.py @@ -30,11 +30,7 @@ from typing_extensions import assert_never from nerfstudio.cameras.lie_groups import exp_map_SE3, exp_map_SO3xR3 from nerfstudio.configs.base_config import InstantiateConfig -from nerfstudio.engine.optimizers import AdamOptimizerConfig, OptimizerConfig -from nerfstudio.engine.schedulers import ( - ExponentialDecaySchedulerConfig, - SchedulerConfig, -) +from nerfstudio.cameras.rays import RayBundle from nerfstudio.utils import poses as pose_utils @@ -47,22 +43,6 @@ class CameraOptimizerConfig(InstantiateConfig): mode: Literal["off", "SO3xR3", "SE3"] = "off" """Pose optimization strategy to use. If enabled, we recommend SO3xR3.""" - position_noise_std: float = 0.0 - """Noise to add to initial positions. Useful for debugging.""" - - orientation_noise_std: float = 0.0 - """Noise to add to initial orientations. Useful for debugging.""" - - optimizer: OptimizerConfig = AdamOptimizerConfig(lr=6e-4, eps=1e-15) - """ADAM parameters for camera optimization.""" - - scheduler: SchedulerConfig = ExponentialDecaySchedulerConfig(max_steps=10000) - """Learning rate scheduler for camera optimizer..""" - - param_group: tyro.conf.Suppress[str] = "camera_opt" - """Name of the parameter group used for pose optimization. Can be any string that doesn't conflict with other - groups.""" - class CameraOptimizer(nn.Module): """Layer that modifies camera poses to be optimized as well as the field during training.""" @@ -89,16 +69,6 @@ class CameraOptimizer(nn.Module): else: assert_never(self.config.mode) - # Initialize pose noise; useful for debugging. - if config.position_noise_std != 0.0 or config.orientation_noise_std != 0.0: - assert config.position_noise_std >= 0.0 and config.orientation_noise_std >= 0.0 - std_vector = torch.tensor( - [config.position_noise_std] * 3 + [config.orientation_noise_std] * 3, device=device - ) - self.pose_noise = exp_map_SE3(torch.normal(torch.zeros((num_cameras, 6), device=device), std_vector)) - else: - self.pose_noise = None - def forward( self, indices: Int[Tensor, "num_cameras"], @@ -122,12 +92,16 @@ class CameraOptimizer(nn.Module): else: assert_never(self.config.mode) - # Apply initial pose noise. - if self.pose_noise is not None: - outputs.append(self.pose_noise[indices, :, :]) - # Return: identity if no transforms are needed, otherwise multiply transforms together. if len(outputs) == 0: # Note that using repeat() instead of tile() here would result in unnecessary copies. return torch.eye(4, device=self.device)[None, :3, :4].tile(indices.shape[0], 1, 1) return functools.reduce(pose_utils.multiply, outputs) + + def apply_to_raybundle(self, raybundle: RayBundle) -> None: + """ + Apply the pose correction to the raybundle + """ + correction_matrices = self(raybundle.camera_indices.squeeze()) + raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3] + raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze() diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index e0a270ba9..96f7c4ccf 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -87,486 +87,492 @@ method_configs["nerfacto"] = TrainerConfig( dataparser=NerfstudioDataParserConfig(), 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), - scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), - ), - ), - model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), - ), - optimizers={ - "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), - }, - "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) -method_configs["nerfacto-big"] = TrainerConfig( - method_name="nerfacto", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=100000, - mixed_precision=True, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=8192, - eval_num_rays_per_batch=4096, - camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", - optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3), - ), ), model=NerfactoModelConfig( eval_num_rays_per_chunk=1 << 15, - num_nerf_samples_per_ray=128, - num_proposal_samples_per_ray=(512, 256), - hidden_dim=128, - hidden_dim_color=128, - appearance_embed_dim=128, - base_res=32, - max_res=4096, - proposal_weights_anneal_max_num_iters=5000, - log2_hashmap_size=21, - ), - ), - optimizers={ - "proposal_networks": { - "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - "fields": { - "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) -method_configs["nerfacto-huge"] = TrainerConfig( - method_name="nerfacto", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=100000, - mixed_precision=True, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=16384, - eval_num_rays_per_batch=4096, - camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", - optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3), - scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000), - ), - ), - model=NerfactoModelConfig( - eval_num_rays_per_chunk=1 << 15, - num_nerf_samples_per_ray=64, - num_proposal_samples_per_ray=(512, 512), - proposal_net_args_list=[ - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False}, - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 7, "max_res": 2048, "use_linear": False}, - ], - hidden_dim=256, - hidden_dim_color=256, - appearance_embed_dim=32, - features_per_level=4, - base_res=32, - max_res=8192, - proposal_weights_anneal_max_num_iters=5000, - log2_hashmap_size=21, - ), - ), - optimizers={ - "proposal_networks": { - "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - "fields": { - "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) - -method_configs["depth-nerfacto"] = TrainerConfig( - method_name="depth-nerfacto", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=True, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - _target=VanillaDataManager[DepthDataset], - dataparser=NerfstudioDataParserConfig(), - 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=DepthNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), - ), - optimizers={ - "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) - -method_configs["volinga"] = TrainerConfig( - method_name="volinga", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=True, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=NerfstudioDataParserConfig(), - 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=NerfactoModelConfig( - eval_num_rays_per_chunk=1 << 15, - hidden_dim=32, - hidden_dim_color=32, - hidden_dim_transient=32, - num_nerf_samples_per_ray=24, - proposal_net_args_list=[ - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": True}, - {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": True}, - ], + camera_optimizer=CameraOptimizerConfig(mode="SO3xR3"), ), ), optimizers={ "proposal_networks": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, + "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), }, - "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) - -method_configs["instant-ngp"] = TrainerConfig( - method_name="instant-ngp", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=True, - pipeline=DynamicBatchPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=NerfstudioDataParserConfig(), - train_num_rays_per_batch=4096, - eval_num_rays_per_batch=4096, - ), - model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), - ), - optimizers={ "fields": { "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), - } - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 12), - vis="viewer", -) - - -method_configs["instant-ngp-bounded"] = TrainerConfig( - method_name="instant-ngp-bounded", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=True, - pipeline=DynamicBatchPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192), - model=InstantNGPModelConfig( - eval_num_rays_per_chunk=8192, - grid_levels=1, - alpha_thre=0.0, - cone_angle=0.0, - disable_scene_contraction=True, - near_plane=0.01, - background_color="black", - ), - ), - optimizers={ - "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), - } - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 12), - vis="viewer", -) - - -method_configs["mipnerf"] = TrainerConfig( - method_name="mipnerf", - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024), - model=VanillaModelConfig( - _target=MipNerfModel, - loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0}, - num_coarse_samples=128, - num_importance_samples=128, - eval_num_rays_per_chunk=1024, - ), - ), - optimizers={ - "fields": { - "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), - "scheduler": None, - } - }, -) - -method_configs["semantic-nerfw"] = TrainerConfig( - method_name="semantic-nerfw", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=True, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - _target=VanillaDataManager[SemanticDataset], - dataparser=Sitcoms3DDataParserConfig(), - train_num_rays_per_batch=4096, - eval_num_rays_per_batch=8192, - ), - model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16), - ), - optimizers={ - "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, }, - "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 16), - vis="viewer", -) - -method_configs["vanilla-nerf"] = TrainerConfig( - method_name="vanilla-nerf", - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=BlenderDataParserConfig(), - ), - model=VanillaModelConfig(_target=NeRFModel), - ), - optimizers={ - "fields": { - "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), - "scheduler": None, - }, - "temporal_distortion": { - "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), - "scheduler": None, - }, - }, -) - -method_configs["tensorf"] = TrainerConfig( - method_name="tensorf", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=False, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=BlenderDataParserConfig(), - train_num_rays_per_batch=4096, - eval_num_rays_per_batch=4096, - ), - model=TensoRFModelConfig( - regularization="tv", - ), - ), - optimizers={ - "fields": { - "optimizer": AdamOptimizerConfig(lr=0.001), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000), - }, - "encodings": { - "optimizer": AdamOptimizerConfig(lr=0.02), - "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000), + "camera_opt": { + "optimizer": AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2), + "scheduler": ExponentialDecaySchedulerConfig(lr_final=6e-6, max_steps=200000), }, }, viewer=ViewerConfig(num_rays_per_chunk=1 << 15), vis="viewer", ) -method_configs["dnerf"] = TrainerConfig( - method_name="dnerf", - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()), - model=VanillaModelConfig( - _target=NeRFModel, - enable_temporal_distortion=True, - temporal_distortion_params={"kind": TemporalDistortionKind.DNERF}, - ), - ), - optimizers={ - "fields": { - "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), - "scheduler": None, - }, - "temporal_distortion": { - "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), - "scheduler": None, - }, - }, -) +# **For now I commented everything out since the datamanager config breaks ** -method_configs["phototourism"] = TrainerConfig( - method_name="phototourism", - steps_per_eval_batch=500, - steps_per_save=2000, - max_num_iterations=30000, - mixed_precision=True, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto - 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) - ), - # Large dataset, so using prior values from VariableResDataManager. - train_num_images_to_sample_from=40, - train_num_times_to_repeat_images=100, - eval_num_images_to_sample_from=40, - eval_num_times_to_repeat_images=100, - ), - model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), - ), - optimizers={ - "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - "fields": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": None, - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) -method_configs["neus"] = TrainerConfig( - method_name="neus", - steps_per_eval_image=500, - steps_per_eval_batch=5000, - steps_per_save=20000, - steps_per_eval_all_images=1000000, # set to a very large number so we don't eval with all images - max_num_iterations=100000, - mixed_precision=False, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - _target=VanillaDataManager[SDFDataset], - dataparser=SDFStudioDataParserConfig(), - train_num_rays_per_batch=1024, - eval_num_rays_per_batch=1024, - camera_optimizer=CameraOptimizerConfig( - mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) - ), - ), - model=NeuSModelConfig(eval_num_rays_per_chunk=1024), - ), - optimizers={ - "fields": { - "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), - "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000), - }, - "field_background": { - "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), - "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000), - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) +# method_configs["nerfacto-big"] = TrainerConfig( +# method_name="nerfacto", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=100000, +# mixed_precision=True, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=NerfstudioDataParserConfig(), +# train_num_rays_per_batch=8192, +# eval_num_rays_per_batch=4096, +# camera_optimizer=CameraOptimizerConfig( +# mode="SO3xR3", +# optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3), +# ), +# ), +# model=NerfactoModelConfig( +# eval_num_rays_per_chunk=1 << 15, +# num_nerf_samples_per_ray=128, +# num_proposal_samples_per_ray=(512, 256), +# hidden_dim=128, +# hidden_dim_color=128, +# appearance_embed_dim=128, +# base_res=32, +# max_res=4096, +# proposal_weights_anneal_max_num_iters=5000, +# log2_hashmap_size=21, +# ), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# "fields": { +# "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) +# method_configs["nerfacto-huge"] = TrainerConfig( +# method_name="nerfacto", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=100000, +# mixed_precision=True, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=NerfstudioDataParserConfig(), +# train_num_rays_per_batch=16384, +# eval_num_rays_per_batch=4096, +# camera_optimizer=CameraOptimizerConfig( +# mode="SO3xR3", +# optimizer=RAdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-3), +# scheduler=ExponentialDecaySchedulerConfig(lr_final=6e-5, max_steps=50000), +# ), +# ), +# model=NerfactoModelConfig( +# eval_num_rays_per_chunk=1 << 15, +# num_nerf_samples_per_ray=64, +# num_proposal_samples_per_ray=(512, 512), +# proposal_net_args_list=[ +# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 512, "use_linear": False}, +# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 7, "max_res": 2048, "use_linear": False}, +# ], +# hidden_dim=256, +# hidden_dim_color=256, +# appearance_embed_dim=32, +# features_per_level=4, +# base_res=32, +# max_res=8192, +# proposal_weights_anneal_max_num_iters=5000, +# log2_hashmap_size=21, +# ), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# "fields": { +# "optimizer": RAdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": ExponentialDecaySchedulerConfig(lr_final=1e-4, max_steps=50000), +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) -method_configs["neus-facto"] = TrainerConfig( - method_name="neus-facto", - steps_per_eval_image=5000, - steps_per_eval_batch=5000, - steps_per_save=2000, - steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images - max_num_iterations=20001, - mixed_precision=False, - pipeline=VanillaPipelineConfig( - datamanager=VanillaDataManagerConfig( - _target=VanillaDataManager[SDFDataset], - dataparser=SDFStudioDataParserConfig(), - train_num_rays_per_batch=2048, - eval_num_rays_per_batch=2048, - camera_optimizer=CameraOptimizerConfig( - mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) - ), - ), - model=NeuSFactoModelConfig( - # proposal network allows for significantly smaller sdf/color network - sdf_field=SDFFieldConfig( - use_grid_feature=True, - num_layers=2, - num_layers_color=2, - hidden_dim=256, - bias=0.5, - beta_init=0.8, - use_appearance_embedding=False, - ), - background_model="none", - eval_num_rays_per_chunk=2048, - ), - ), - optimizers={ - "proposal_networks": { - "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), - "scheduler": MultiStepSchedulerConfig(max_steps=20001, milestones=(10000, 1500, 18000)), - }, - "fields": { - "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), - "scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001), - }, - "field_background": { - "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), - "scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001), - }, - }, - viewer=ViewerConfig(num_rays_per_chunk=1 << 15), - vis="viewer", -) +# method_configs["depth-nerfacto"] = TrainerConfig( +# method_name="depth-nerfacto", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=True, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# _target=VanillaDataManager[DepthDataset], +# dataparser=NerfstudioDataParserConfig(), +# 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=DepthNerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) + +# method_configs["volinga"] = TrainerConfig( +# method_name="volinga", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=True, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=NerfstudioDataParserConfig(), +# 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=NerfactoModelConfig( +# eval_num_rays_per_chunk=1 << 15, +# hidden_dim=32, +# hidden_dim_color=32, +# hidden_dim_transient=32, +# num_nerf_samples_per_ray=24, +# proposal_net_args_list=[ +# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 128, "use_linear": True}, +# {"hidden_dim": 16, "log2_hashmap_size": 17, "num_levels": 5, "max_res": 256, "use_linear": True}, +# ], +# ), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) + +# method_configs["instant-ngp"] = TrainerConfig( +# method_name="instant-ngp", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=True, +# pipeline=DynamicBatchPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=NerfstudioDataParserConfig(), +# train_num_rays_per_batch=4096, +# eval_num_rays_per_batch=4096, +# ), +# model=InstantNGPModelConfig(eval_num_rays_per_chunk=8192), +# ), +# optimizers={ +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), +# } +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 12), +# vis="viewer", +# ) + + +# method_configs["instant-ngp-bounded"] = TrainerConfig( +# method_name="instant-ngp-bounded", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=True, +# pipeline=DynamicBatchPipelineConfig( +# datamanager=VanillaDataManagerConfig(dataparser=InstantNGPDataParserConfig(), train_num_rays_per_batch=8192), +# model=InstantNGPModelConfig( +# eval_num_rays_per_chunk=8192, +# grid_levels=1, +# alpha_thre=0.0, +# cone_angle=0.0, +# disable_scene_contraction=True, +# near_plane=0.01, +# background_color="black", +# ), +# ), +# optimizers={ +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=200000), +# } +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 12), +# vis="viewer", +# ) + + +# method_configs["mipnerf"] = TrainerConfig( +# method_name="mipnerf", +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig(dataparser=NerfstudioDataParserConfig(), train_num_rays_per_batch=1024), +# model=VanillaModelConfig( +# _target=MipNerfModel, +# loss_coefficients={"rgb_loss_coarse": 0.1, "rgb_loss_fine": 1.0}, +# num_coarse_samples=128, +# num_importance_samples=128, +# eval_num_rays_per_chunk=1024, +# ), +# ), +# optimizers={ +# "fields": { +# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), +# "scheduler": None, +# } +# }, +# ) + +# method_configs["semantic-nerfw"] = TrainerConfig( +# method_name="semantic-nerfw", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=True, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# _target=VanillaDataManager[SemanticDataset], +# dataparser=Sitcoms3DDataParserConfig(), +# train_num_rays_per_batch=4096, +# eval_num_rays_per_batch=8192, +# ), +# model=SemanticNerfWModelConfig(eval_num_rays_per_chunk=1 << 16), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 16), +# vis="viewer", +# ) + +# method_configs["vanilla-nerf"] = TrainerConfig( +# method_name="vanilla-nerf", +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=BlenderDataParserConfig(), +# ), +# model=VanillaModelConfig(_target=NeRFModel), +# ), +# optimizers={ +# "fields": { +# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), +# "scheduler": None, +# }, +# "temporal_distortion": { +# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), +# "scheduler": None, +# }, +# }, +# ) + +# method_configs["tensorf"] = TrainerConfig( +# method_name="tensorf", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=False, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=BlenderDataParserConfig(), +# train_num_rays_per_batch=4096, +# eval_num_rays_per_batch=4096, +# ), +# model=TensoRFModelConfig( +# regularization="tv", +# ), +# ), +# optimizers={ +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=0.001), +# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.0001, max_steps=30000), +# }, +# "encodings": { +# "optimizer": AdamOptimizerConfig(lr=0.02), +# "scheduler": ExponentialDecaySchedulerConfig(lr_final=0.002, max_steps=30000), +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) + +# method_configs["dnerf"] = TrainerConfig( +# method_name="dnerf", +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig(dataparser=DNeRFDataParserConfig()), +# model=VanillaModelConfig( +# _target=NeRFModel, +# enable_temporal_distortion=True, +# temporal_distortion_params={"kind": TemporalDistortionKind.DNERF}, +# ), +# ), +# optimizers={ +# "fields": { +# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), +# "scheduler": None, +# }, +# "temporal_distortion": { +# "optimizer": RAdamOptimizerConfig(lr=5e-4, eps=1e-08), +# "scheduler": None, +# }, +# }, +# ) + +# method_configs["phototourism"] = TrainerConfig( +# method_name="phototourism", +# steps_per_eval_batch=500, +# steps_per_save=2000, +# max_num_iterations=30000, +# mixed_precision=True, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# dataparser=PhototourismDataParserConfig(), # NOTE: one of the only differences with nerfacto +# 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) +# ), +# # Large dataset, so using prior values from VariableResDataManager. +# train_num_images_to_sample_from=40, +# train_num_times_to_repeat_images=100, +# eval_num_images_to_sample_from=40, +# eval_num_times_to_repeat_images=100, +# ), +# model=NerfactoModelConfig(eval_num_rays_per_chunk=1 << 15), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": None, +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) + +# method_configs["neus"] = TrainerConfig( +# method_name="neus", +# steps_per_eval_image=500, +# steps_per_eval_batch=5000, +# steps_per_save=20000, +# steps_per_eval_all_images=1000000, # set to a very large number so we don't eval with all images +# max_num_iterations=100000, +# mixed_precision=False, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# _target=VanillaDataManager[SDFDataset], +# dataparser=SDFStudioDataParserConfig(), +# train_num_rays_per_batch=1024, +# eval_num_rays_per_batch=1024, +# camera_optimizer=CameraOptimizerConfig( +# mode="off", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) +# ), +# ), +# model=NeuSModelConfig(eval_num_rays_per_chunk=1024), +# ), +# optimizers={ +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), +# "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000), +# }, +# "field_background": { +# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), +# "scheduler": CosineDecaySchedulerConfig(warm_up_end=5000, learning_rate_alpha=0.05, max_steps=300000), +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) + +# method_configs["neus-facto"] = TrainerConfig( +# method_name="neus-facto", +# steps_per_eval_image=5000, +# steps_per_eval_batch=5000, +# steps_per_save=2000, +# steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images +# max_num_iterations=20001, +# mixed_precision=False, +# pipeline=VanillaPipelineConfig( +# datamanager=VanillaDataManagerConfig( +# _target=VanillaDataManager[SDFDataset], +# dataparser=SDFStudioDataParserConfig(), +# train_num_rays_per_batch=2048, +# eval_num_rays_per_batch=2048, +# camera_optimizer=CameraOptimizerConfig( +# mode="SO3xR3", optimizer=AdamOptimizerConfig(lr=6e-4, eps=1e-8, weight_decay=1e-2) +# ), +# ), +# model=NeuSFactoModelConfig( +# # proposal network allows for significantly smaller sdf/color network +# sdf_field=SDFFieldConfig( +# use_grid_feature=True, +# num_layers=2, +# num_layers_color=2, +# hidden_dim=256, +# bias=0.5, +# beta_init=0.8, +# use_appearance_embedding=False, +# ), +# background_model="none", +# eval_num_rays_per_chunk=2048, +# ), +# ), +# optimizers={ +# "proposal_networks": { +# "optimizer": AdamOptimizerConfig(lr=1e-2, eps=1e-15), +# "scheduler": MultiStepSchedulerConfig(max_steps=20001, milestones=(10000, 1500, 18000)), +# }, +# "fields": { +# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), +# "scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001), +# }, +# "field_background": { +# "optimizer": AdamOptimizerConfig(lr=5e-4, eps=1e-15), +# "scheduler": CosineDecaySchedulerConfig(warm_up_end=500, learning_rate_alpha=0.05, max_steps=20001), +# }, +# }, +# viewer=ViewerConfig(num_rays_per_chunk=1 << 15), +# vis="viewer", +# ) def merge_methods(methods, method_descriptions, new_methods, new_descriptions, overwrite=True): diff --git a/nerfstudio/data/datamanagers/base_datamanager.py b/nerfstudio/data/datamanagers/base_datamanager.py index 82e4f66d3..c957af0ca 100644 --- a/nerfstudio/data/datamanagers/base_datamanager.py +++ b/nerfstudio/data/datamanagers/base_datamanager.py @@ -43,7 +43,6 @@ from torch.nn import Parameter from torch.utils.data.distributed import DistributedSampler from typing_extensions import TypeVar -from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig from nerfstudio.cameras.cameras import CameraType from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.base_config import InstantiateConfig @@ -109,8 +108,6 @@ class DataManagerConfig(InstantiateConfig): """Target class to instantiate.""" data: Optional[Path] = None """Source of data, may not be used by all models.""" - camera_optimizer: Optional[CameraOptimizerConfig] = None - """Specifies the camera pose optimizer used during training. Helpful if poses are noisy.""" class DataManager(nn.Module): @@ -329,9 +326,6 @@ class VanillaDataManagerConfig(DataManagerConfig): new images. If -1, never pick new images.""" eval_image_indices: Optional[Tuple[int, ...]] = (0,) """Specifies the image indices to use during eval; if None, uses all.""" - 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] = cast(Any, staticmethod(nerfstudio_collate)) """Specifies the collate function to use for the train and eval dataloaders.""" camera_res_scale_factor: float = 1.0 @@ -459,13 +453,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]): ) self.iter_train_image_dataloader = iter(self.train_image_dataloader) self.train_pixel_sampler = self._get_pixel_sampler(self.train_dataset, self.config.train_num_rays_per_batch) - self.train_camera_optimizer = self.config.camera_optimizer.setup( - num_cameras=self.train_dataset.cameras.size, device=self.device - ) - self.train_ray_generator = RayGenerator( - self.train_dataset.cameras.to(self.device), - self.train_camera_optimizer, - ) + self.train_ray_generator = RayGenerator(self.train_dataset.cameras.to(self.device)) def setup_eval(self): """Sets up the data loader for evaluation""" @@ -483,13 +471,7 @@ class VanillaDataManager(DataManager, Generic[TDataset]): ) self.iter_eval_image_dataloader = iter(self.eval_image_dataloader) self.eval_pixel_sampler = self._get_pixel_sampler(self.eval_dataset, self.config.eval_num_rays_per_batch) - self.eval_camera_optimizer = self.config.camera_optimizer.setup( - num_cameras=self.eval_dataset.cameras.size, device=self.device - ) - self.eval_ray_generator = RayGenerator( - self.eval_dataset.cameras.to(self.device), - self.eval_camera_optimizer, - ) + self.eval_ray_generator = RayGenerator(self.eval_dataset.cameras.to(self.device)) # for loading full images self.fixed_indices_eval_dataloader = FixedIndicesEvalDataloader( input_dataset=self.eval_dataset, @@ -545,13 +527,4 @@ class VanillaDataManager(DataManager, Generic[TDataset]): Returns: A list of dictionaries containing the data manager's param groups. """ - param_groups = {} - - camera_opt_params = list(self.train_camera_optimizer.parameters()) - if self.config.camera_optimizer.mode != "off": - assert len(camera_opt_params) > 0 - param_groups[self.config.camera_optimizer.param_group] = camera_opt_params - else: - assert len(camera_opt_params) == 0 - - return param_groups + return {} diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index afc0d94f3..805de6cd9 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -211,13 +211,6 @@ class Trainer: """ optimizer_config = self.config.optimizers.copy() param_groups = self.pipeline.get_param_groups() - camera_optimizer_config = self.config.pipeline.datamanager.camera_optimizer - if camera_optimizer_config is not None and camera_optimizer_config.mode != "off": - assert camera_optimizer_config.param_group not in optimizer_config - optimizer_config[camera_optimizer_config.param_group] = { - "optimizer": camera_optimizer_config.optimizer, - "scheduler": camera_optimizer_config.scheduler, - } return Optimizers(optimizer_config, param_groups) def train(self) -> None: diff --git a/nerfstudio/model_components/ray_generators.py b/nerfstudio/model_components/ray_generators.py index 49754548d..bad32258d 100644 --- a/nerfstudio/model_components/ray_generators.py +++ b/nerfstudio/model_components/ray_generators.py @@ -18,7 +18,6 @@ Ray generator. from jaxtyping import Int from torch import Tensor, nn -from nerfstudio.cameras.camera_optimizers import CameraOptimizer from nerfstudio.cameras.cameras import Cameras from nerfstudio.cameras.rays import RayBundle @@ -34,10 +33,9 @@ class RayGenerator(nn.Module): image_coords: Tensor - def __init__(self, cameras: Cameras, pose_optimizer: CameraOptimizer) -> None: + def __init__(self, cameras: Cameras) -> None: super().__init__() self.cameras = cameras - self.pose_optimizer = pose_optimizer self.register_buffer("image_coords", cameras.get_image_coords(), persistent=False) def forward(self, ray_indices: Int[Tensor, "num_rays 3"]) -> RayBundle: @@ -51,11 +49,8 @@ class RayGenerator(nn.Module): x = ray_indices[:, 2] # col indices coords = self.image_coords[y, x] - camera_opt_to_camera = self.pose_optimizer(c) - ray_bundle = self.cameras.generate_rays( camera_indices=c.unsqueeze(-1), coords=coords, - camera_opt_to_camera=camera_opt_to_camera, ) return ray_bundle diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 75e798873..2bbfc51e2 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -48,6 +48,7 @@ from nerfstudio.model_components.scene_colliders import NearFarCollider from nerfstudio.model_components.shaders import NormalsShader from nerfstudio.models.base_model import Model, ModelConfig from nerfstudio.utils import colormaps +from nerfstudio.cameras.camera_optimizers import CameraOptimizerConfig, CameraOptimizer @dataclass @@ -126,6 +127,8 @@ class NerfactoModelConfig(ModelConfig): """Which implementation to use for the model.""" appearance_embed_dim: int = 32 """Dimension of the appearance embedding.""" + camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig() + """Config of the camera optimizer to use""" class NerfactoModel(Model): @@ -163,6 +166,9 @@ class NerfactoModel(Model): implementation=self.config.implementation, ) + self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup( + num_cameras=self.num_train_data, device="cpu" # for now initialize on cpu, these will be moved later + ) self.density_fns = [] num_prop_nets = self.config.num_proposal_iterations # Build the proposal network(s) @@ -236,6 +242,12 @@ class NerfactoModel(Model): param_groups = {} param_groups["proposal_networks"] = list(self.proposal_networks.parameters()) param_groups["fields"] = list(self.field.parameters()) + camera_opt_params = list(self.camera_optimizer.parameters()) + if self.config.camera_optimizer.mode != "off": + assert len(camera_opt_params) > 0 + param_groups["camera_opt"] = camera_opt_params + else: + assert len(camera_opt_params) == 0 return param_groups def get_training_callbacks( @@ -273,6 +285,10 @@ class NerfactoModel(Model): return callbacks def get_outputs(self, ray_bundle: RayBundle): + # apply the camera optimizer pose tweaks + if self.training: + self.camera_optimizer.apply_to_raybundle(ray_bundle) + ray_samples: RaySamples ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns) field_outputs = self.field.forward(ray_samples, compute_normals=self.config.predict_normals) @@ -325,6 +341,10 @@ class NerfactoModel(Model): metrics_dict["psnr"] = self.psnr(outputs["rgb"], image) if self.training: metrics_dict["distortion"] = distortion_loss(outputs["weights_list"], outputs["ray_samples_list"]) + + # Report the camera optimization metrics + metrics_dict["camera_opt_translation"] = self.camera_optimizer.pose_adjustment[:, :3].norm() + metrics_dict["camera_opt_rotation"] = self.camera_optimizer.pose_adjustment[:, 3:].norm() return metrics_dict def get_loss_dict(self, outputs, batch, metrics_dict=None): diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index 1e1e0fab5..a3a6006ab 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -280,18 +280,6 @@ class VanillaPipeline(Pipeline): ray_bundle, batch = self.datamanager.next_train(step) model_outputs = self._model(ray_bundle) # train distributed data parallel model if world_size > 1 metrics_dict = self.model.get_metrics_dict(model_outputs, batch) - - if self.config.datamanager.camera_optimizer is not None: - camera_opt_param_group = self.config.datamanager.camera_optimizer.param_group - if camera_opt_param_group in self.datamanager.get_param_groups(): - # Report the camera optimization metrics - metrics_dict["camera_opt_translation"] = ( - self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, :3].norm() - ) - metrics_dict["camera_opt_rotation"] = ( - self.datamanager.get_param_groups()[camera_opt_param_group][0].data[:, 3:].norm() - ) - loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) return model_outputs, loss_dict, metrics_dict