Spaces:
Running
Running
File size: 6,027 Bytes
7088d16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 |
# 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.renderer.implicit import HarmonicEmbedding
from torch.distributions import MultivariateNormal
from .common_testing import TestCaseMixin
class TestHarmonicEmbedding(TestCaseMixin, unittest.TestCase):
def setUp(self) -> None:
super().setUp()
torch.manual_seed(1)
def test_correct_output_dim(self):
embed_fun = HarmonicEmbedding(n_harmonic_functions=2, append_input=False)
# input_dims * (2 * n_harmonic_functions + int(append_input))
output_dim = 3 * (2 * 2 + int(False))
self.assertEqual(
output_dim,
embed_fun.get_output_dim_static(
input_dims=3, n_harmonic_functions=2, append_input=False
),
)
self.assertEqual(output_dim, embed_fun.get_output_dim())
def test_correct_frequency_range(self):
embed_fun_log = HarmonicEmbedding(n_harmonic_functions=3)
embed_fun_lin = HarmonicEmbedding(n_harmonic_functions=3, logspace=False)
self.assertClose(embed_fun_log._frequencies, torch.FloatTensor((1.0, 2.0, 4.0)))
self.assertClose(embed_fun_lin._frequencies, torch.FloatTensor((1.0, 2.5, 4.0)))
def test_correct_embed_out(self):
n_harmonic_functions = 2
x = torch.randn((1, 5))
D = 5 * n_harmonic_functions * 2 # sin + cos
embed_fun = HarmonicEmbedding(
n_harmonic_functions=n_harmonic_functions, append_input=False
)
embed_out = embed_fun(x)
self.assertEqual(embed_out.shape, (1, D))
# Sum the squares of the respective frequencies
# cos^2(x) + sin^2(x) = 1
sum_squares = embed_out[0, : D // 2] ** 2 + embed_out[0, D // 2 :] ** 2
self.assertClose(sum_squares, torch.ones((D // 2)))
# Test append input
embed_fun = HarmonicEmbedding(
n_harmonic_functions=n_harmonic_functions, append_input=True
)
embed_out_appended_input = embed_fun(x)
self.assertClose(
embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
)
# Last plane in output is the input
self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)
def test_correct_embed_out_with_diag_cov(self):
n_harmonic_functions = 2
x = torch.randn((1, 3))
diag_cov = torch.randn((1, 3))
D = 3 * n_harmonic_functions * 2 # sin + cos
embed_fun = HarmonicEmbedding(
n_harmonic_functions=n_harmonic_functions, append_input=False
)
embed_out = embed_fun(x, diag_cov=diag_cov)
self.assertEqual(embed_out.shape, (1, D))
# Compute the scaling factor introduce in MipNerf
scale_factor = (
-0.5 * diag_cov[..., None] * torch.pow(embed_fun._frequencies[None, :], 2)
)
scale_factor = torch.exp(scale_factor).reshape(1, -1).tile((1, 2))
# If we remove this scaling factor, we should go back to the
# classical harmonic embedding:
# Sum the squares of the respective frequencies
# cos^2(x) + sin^2(x) = 1
embed_out_without_cov = embed_out / scale_factor
sum_squares = (
embed_out_without_cov[0, : D // 2] ** 2
+ embed_out_without_cov[0, D // 2 :] ** 2
)
self.assertClose(sum_squares, torch.ones((D // 2)))
# Test append input
embed_fun = HarmonicEmbedding(
n_harmonic_functions=n_harmonic_functions, append_input=True
)
embed_out_appended_input = embed_fun(x, diag_cov=diag_cov)
self.assertClose(
embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
)
# Last plane in output is the input
self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)
def test_correct_behavior_between_ipe_and_its_estimation_from_harmonic_embedding(
self,
):
"""
Check that the HarmonicEmbedding with integrated_position_encoding (IPE) set to
True is coherent with the HarmonicEmbedding.
What is the idea behind this test?
We wish to produce an IPE that is the expectation
of our lifted multivariate gaussian, modulated by the sine and cosine of
the coordinates. These expectation has a closed-form
(see equations 11, 12, 13, 14 of [1]).
We sample N elements from the multivariate gaussian defined by its mean and covariance
and compute the HarmonicEmbedding. The expected value of those embeddings should be
equal to our IPE.
Inspired from:
https://github.com/google/mipnerf/blob/84c969e0a623edd183b75693aed72a7e7c22902d/internal/mip_test.py#L359
References:
[1] `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_.
"""
num_dims = 3
n_harmonic_functions = 6
mean = torch.randn(num_dims)
diag_cov = torch.rand(num_dims)
he_fun = HarmonicEmbedding(
n_harmonic_functions=n_harmonic_functions, logspace=True, append_input=False
)
ipe_fun = HarmonicEmbedding(
n_harmonic_functions=n_harmonic_functions,
append_input=False,
)
embedding_ipe = ipe_fun(mean, diag_cov=diag_cov)
rand_mvn = MultivariateNormal(mean, torch.eye(num_dims) * diag_cov)
# Providing a large enough number of samples
# we should obtain an estimation close to our IPE
num_samples = 100000
embedding_he = he_fun(rand_mvn.sample_n(num_samples))
self.assertClose(embedding_he.mean(0), embedding_ipe, rtol=1e-2, atol=1e-2)
|