mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-01 03:12:49 +08:00
813e941de5
Summary: Introduces the OverfitModel for NeRF-style training with overfitting to one scene. It is a specific case of GenericModel. It has been disentangle to ease usage. ## General modification 1. Modularize a minimum GenericModel to introduce OverfitModel 2. Introduce OverfitModel and ensure through unit testing that it behaves like GenericModel. ## Modularization The following methods have been extracted from GenericModel to allow modularity with ManyViewModel: - get_objective is now a call to weighted_sum_losses - log_loss_weights - prepare_inputs The generic methods have been moved to an utils.py file. Simplify the code to introduce OverfitModel. Private methods like chunk_generator are now public and can now be used by ManyViewModel. Reviewed By: shapovalov Differential Revision: D43771992 fbshipit-source-id: 6102aeb21c7fdd56aa2ff9cd1dd23fd9fbf26315
67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
# 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
|
|
|
|
import torch
|
|
|
|
from pytorch3d.implicitron.models.utils import preprocess_input, weighted_sum_losses
|
|
|
|
|
|
class TestUtils(unittest.TestCase):
|
|
def test_prepare_inputs_wrong_num_dim(self):
|
|
img = torch.randn(3, 3, 3)
|
|
with self.assertRaises(ValueError) as context:
|
|
img, fg_prob, depth_map = preprocess_input(
|
|
img, None, None, True, True, 0.5, (0.0, 0.0, 0.0)
|
|
)
|
|
self.assertEqual(
|
|
"Model received unbatched inputs. "
|
|
+ "Perhaps they came from a FrameData which had not been collated.",
|
|
context.exception,
|
|
)
|
|
|
|
def test_prepare_inputs_mask_image_true(self):
|
|
batch, channels, height, width = 2, 3, 10, 10
|
|
img = torch.ones(batch, channels, height, width)
|
|
# Create a mask on the lower triangular matrix
|
|
fg_prob = torch.tril(torch.ones(batch, 1, height, width)) * 0.3
|
|
|
|
out_img, out_fg_prob, out_depth_map = preprocess_input(
|
|
img, fg_prob, None, True, False, 0.3, (0.0, 0.0, 0.0)
|
|
)
|
|
|
|
self.assertTrue(torch.equal(out_img, torch.tril(img)))
|
|
self.assertTrue(torch.equal(out_fg_prob, fg_prob >= 0.3))
|
|
self.assertIsNone(out_depth_map)
|
|
|
|
def test_prepare_inputs_mask_depth_true(self):
|
|
batch, channels, height, width = 2, 3, 10, 10
|
|
img = torch.ones(batch, channels, height, width)
|
|
depth_map = torch.randn(batch, channels, height, width)
|
|
# Create a mask on the lower triangular matrix
|
|
fg_prob = torch.tril(torch.ones(batch, 1, height, width)) * 0.3
|
|
|
|
out_img, out_fg_prob, out_depth_map = preprocess_input(
|
|
img, fg_prob, depth_map, False, True, 0.3, (0.0, 0.0, 0.0)
|
|
)
|
|
|
|
self.assertTrue(torch.equal(out_img, img))
|
|
self.assertTrue(torch.equal(out_fg_prob, fg_prob >= 0.3))
|
|
self.assertTrue(torch.equal(out_depth_map, torch.tril(depth_map)))
|
|
|
|
def test_weighted_sum_losses(self):
|
|
preds = {"a": torch.tensor(2), "b": torch.tensor(2)}
|
|
weights = {"a": 2.0, "b": 0.0}
|
|
loss = weighted_sum_losses(preds, weights)
|
|
self.assertEqual(loss, 4.0)
|
|
|
|
def test_weighted_sum_losses_raise_warning(self):
|
|
preds = {"a": torch.tensor(2), "b": torch.tensor(2)}
|
|
weights = {"c": 2.0, "d": 2.0}
|
|
self.assertIsNone(weighted_sum_losses(preds, weights))
|