# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import unittest from typing import Any, Dict from unittest.mock import patch import torch from pytorch3d.implicitron.models.generic_model import GenericModel from pytorch3d.implicitron.models.overfit_model import OverfitModel from pytorch3d.implicitron.models.renderer.base import EvaluationMode from pytorch3d.implicitron.tools.config import expand_args_fields from pytorch3d.renderer.cameras import look_at_view_transform, PerspectiveCameras DEVICE = torch.device("cuda:0") def _generate_fake_inputs(N: int, H: int, W: int) -> Dict[str, Any]: R, T = look_at_view_transform(azim=torch.rand(N) * 360) return { "camera": PerspectiveCameras(R=R, T=T, device=DEVICE), "fg_probability": torch.randint( high=2, size=(N, 1, H, W), device=DEVICE ).float(), "depth_map": torch.rand((N, 1, H, W), device=DEVICE) + 0.1, "mask_crop": torch.randint(high=2, size=(N, 1, H, W), device=DEVICE).float(), "sequence_name": ["sequence"] * N, "image_rgb": torch.rand((N, 1, H, W), device=DEVICE), } def mock_safe_multinomial(input: torch.Tensor, num_samples: int) -> torch.Tensor: """Return non deterministic indexes to mock safe_multinomial Args: input: tensor of shape [B, n] containing non-negative values; rows are interpreted as unnormalized event probabilities in categorical distributions. num_samples: number of samples to take. Returns: Tensor of shape [B, num_samples] """ batch_size = input.shape[0] return torch.arange(num_samples).repeat(batch_size, 1).to(DEVICE) class TestOverfitModel(unittest.TestCase): def setUp(self): torch.manual_seed(42) def test_overfit_model_vs_generic_model_with_batch_size_one(self): """In this test we compare OverfitModel to GenericModel behavior. We use a Nerf setup (2 rendering passes). OverfitModel is a specific case of GenericModel. Hence, with the same inputs, they should provide the exact same results. """ expand_args_fields(OverfitModel) expand_args_fields(GenericModel) batch_size, image_height, image_width = 1, 80, 80 assert batch_size == 1 overfit_model = OverfitModel( render_image_height=image_height, render_image_width=image_width, coarse_implicit_function_class_type="NeuralRadianceFieldImplicitFunction", # To avoid randomization to compare the outputs of our model # we deactivate the stratified_point_sampling_training raysampler_AdaptiveRaySampler_args={ "stratified_point_sampling_training": False }, global_encoder_class_type="SequenceAutodecoder", global_encoder_SequenceAutodecoder_args={ "autodecoder_args": { "n_instances": 1000, "init_scale": 1.0, "encoding_dim": 64, } }, ) generic_model = GenericModel( render_image_height=image_height, render_image_width=image_width, n_train_target_views=batch_size, num_passes=2, # To avoid randomization to compare the outputs of our model # we deactivate the stratified_point_sampling_training raysampler_AdaptiveRaySampler_args={ "stratified_point_sampling_training": False }, global_encoder_class_type="SequenceAutodecoder", global_encoder_SequenceAutodecoder_args={ "autodecoder_args": { "n_instances": 1000, "init_scale": 1.0, "encoding_dim": 64, } }, ) # Check if they do share the number of parameters num_params_mvm = sum(p.numel() for p in overfit_model.parameters()) num_params_gm = sum(p.numel() for p in generic_model.parameters()) self.assertEqual(num_params_mvm, num_params_gm) # Adapt the mapping from generic model to overfit model mapping_om_from_gm = { key.replace( "_implicit_functions.0._fn", "coarse_implicit_function" ).replace("_implicit_functions.1._fn", "implicit_function"): val for key, val in generic_model.state_dict().items() } # Copy parameters from generic_model to overfit_model overfit_model.load_state_dict(mapping_om_from_gm) overfit_model.to(DEVICE) generic_model.to(DEVICE) inputs_ = _generate_fake_inputs(batch_size, image_height, image_width) # training forward pass overfit_model.train() generic_model.train() with patch( "pytorch3d.renderer.implicit.raysampling._safe_multinomial", side_effect=mock_safe_multinomial, ): train_preds_om = overfit_model( **inputs_, evaluation_mode=EvaluationMode.TRAINING, ) train_preds_gm = generic_model( **inputs_, evaluation_mode=EvaluationMode.TRAINING, ) self.assertTrue(len(train_preds_om) == len(train_preds_gm)) self.assertTrue(train_preds_om["objective"].isfinite().item()) # We avoid all the randomization and the weights are the same # The objective should be the same self.assertTrue( torch.allclose(train_preds_om["objective"], train_preds_gm["objective"]) ) # Test if the evaluation works overfit_model.eval() generic_model.eval() with torch.no_grad(): eval_preds_om = overfit_model( **inputs_, evaluation_mode=EvaluationMode.EVALUATION, ) eval_preds_gm = generic_model( **inputs_, evaluation_mode=EvaluationMode.EVALUATION, ) self.assertEqual( eval_preds_om["images_render"].shape, (batch_size, 3, image_height, image_width), ) self.assertTrue( torch.allclose(eval_preds_om["objective"], eval_preds_gm["objective"]) ) self.assertTrue( torch.allclose( eval_preds_om["images_render"], eval_preds_gm["images_render"] ) ) def test_overfit_model_check_share_weights(self): model = OverfitModel(share_implicit_function_across_passes=True) for p1, p2 in zip( model.implicit_function.parameters(), model.coarse_implicit_function.parameters(), ): self.assertEqual(id(p1), id(p2)) model.to(DEVICE) inputs_ = _generate_fake_inputs(2, 80, 80) model(**inputs_, evaluation_mode=EvaluationMode.TRAINING) def test_overfit_model_check_no_share_weights(self): model = OverfitModel( share_implicit_function_across_passes=False, coarse_implicit_function_class_type="NeuralRadianceFieldImplicitFunction", coarse_implicit_function_NeuralRadianceFieldImplicitFunction_args={ "transformer_dim_down_factor": 1.0, "n_hidden_neurons_xyz": 256, "n_layers_xyz": 8, "append_xyz": (5,), }, ) for p1, p2 in zip( model.implicit_function.parameters(), model.coarse_implicit_function.parameters(), ): self.assertNotEqual(id(p1), id(p2)) model.to(DEVICE) inputs_ = _generate_fake_inputs(2, 80, 80) model(**inputs_, evaluation_mode=EvaluationMode.TRAINING) def test_overfit_model_coarse_implicit_function_is_none(self): model = OverfitModel( share_implicit_function_across_passes=False, coarse_implicit_function_NeuralRadianceFieldImplicitFunction_args=None, ) self.assertIsNone(model.coarse_implicit_function) model.to(DEVICE) inputs_ = _generate_fake_inputs(2, 80, 80) model(**inputs_, evaluation_mode=EvaluationMode.TRAINING)