Spaces:
Running
Running
| # 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.ops import packed_to_padded, padded_to_packed | |
| from pytorch3d.structures.meshes import Meshes | |
| from .common_testing import get_random_cuda_device, TestCaseMixin | |
| class TestPackedToPadded(TestCaseMixin, unittest.TestCase): | |
| def setUp(self) -> None: | |
| super().setUp() | |
| torch.manual_seed(1) | |
| def init_meshes( | |
| num_meshes: int = 10, | |
| num_verts: int = 1000, | |
| num_faces: int = 3000, | |
| device: str = "cpu", | |
| ): | |
| device = torch.device(device) | |
| verts_list = [] | |
| faces_list = [] | |
| for _ in range(num_meshes): | |
| verts = torch.rand((num_verts, 3), dtype=torch.float32, device=device) | |
| faces = torch.randint( | |
| num_verts, size=(num_faces, 3), dtype=torch.int64, device=device | |
| ) | |
| verts_list.append(verts) | |
| faces_list.append(faces) | |
| meshes = Meshes(verts_list, faces_list) | |
| return meshes | |
| def packed_to_padded_python(inputs, first_idxs, max_size, device): | |
| """ | |
| PyTorch implementation of packed_to_padded function. | |
| """ | |
| num_meshes = first_idxs.size(0) | |
| if inputs.dim() == 1: | |
| inputs_padded = torch.zeros((num_meshes, max_size), device=device) | |
| else: | |
| inputs_padded = torch.zeros( | |
| (num_meshes, max_size, *inputs.shape[1:]), device=device | |
| ) | |
| for m in range(num_meshes): | |
| s = first_idxs[m] | |
| if m == num_meshes - 1: | |
| f = inputs.shape[0] | |
| else: | |
| f = first_idxs[m + 1] | |
| inputs_padded[m, : f - s] = inputs[s:f] | |
| return inputs_padded | |
| def padded_to_packed_python(inputs, first_idxs, num_inputs, device): | |
| """ | |
| PyTorch implementation of padded_to_packed function. | |
| """ | |
| num_meshes = inputs.size(0) | |
| if inputs.dim() == 2: | |
| inputs_packed = torch.zeros((num_inputs,), device=device) | |
| else: | |
| inputs_packed = torch.zeros((num_inputs, *inputs.shape[2:]), device=device) | |
| for m in range(num_meshes): | |
| s = first_idxs[m] | |
| if m == num_meshes - 1: | |
| f = num_inputs | |
| else: | |
| f = first_idxs[m + 1] | |
| inputs_packed[s:f] = inputs[m, : f - s] | |
| return inputs_packed | |
| def _test_packed_to_padded_helper(self, dims, device): | |
| """ | |
| Check the results from packed_to_padded and PyTorch implementations | |
| are the same. | |
| """ | |
| meshes = self.init_meshes(16, 100, 300, device=device) | |
| faces = meshes.faces_packed() | |
| mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() | |
| max_faces = meshes.num_faces_per_mesh().max().item() | |
| if len(dims) == 0: | |
| values = torch.rand((faces.shape[0],), device=device, requires_grad=True) | |
| else: | |
| values = torch.rand( | |
| (faces.shape[0], *dims), device=device, requires_grad=True | |
| ) | |
| values_torch = values.detach().clone() | |
| values_torch.requires_grad = True | |
| values_padded = packed_to_padded( | |
| values, mesh_to_faces_packed_first_idx, max_faces | |
| ) | |
| values_padded_torch = TestPackedToPadded.packed_to_padded_python( | |
| values_torch, mesh_to_faces_packed_first_idx, max_faces, device | |
| ) | |
| # check forward | |
| self.assertClose(values_padded, values_padded_torch) | |
| # check backward | |
| if len(dims) == 0: | |
| grad_inputs = torch.rand((len(meshes), max_faces), device=device) | |
| else: | |
| grad_inputs = torch.rand((len(meshes), max_faces, *dims), device=device) | |
| values_padded.backward(grad_inputs) | |
| grad_outputs = values.grad | |
| values_padded_torch.backward(grad_inputs) | |
| grad_outputs_torch1 = values_torch.grad | |
| grad_outputs_torch2 = TestPackedToPadded.padded_to_packed_python( | |
| grad_inputs, mesh_to_faces_packed_first_idx, values.size(0), device=device | |
| ) | |
| self.assertClose(grad_outputs, grad_outputs_torch1) | |
| self.assertClose(grad_outputs, grad_outputs_torch2) | |
| def test_packed_to_padded_flat_cpu(self): | |
| self._test_packed_to_padded_helper([], "cpu") | |
| def test_packed_to_padded_D1_cpu(self): | |
| self._test_packed_to_padded_helper([1], "cpu") | |
| def test_packed_to_padded_D16_cpu(self): | |
| self._test_packed_to_padded_helper([16], "cpu") | |
| def test_packed_to_padded_D16_9_cpu(self): | |
| self._test_packed_to_padded_helper([16, 9], "cpu") | |
| def test_packed_to_padded_D16_3_2_cpu(self): | |
| self._test_packed_to_padded_helper([16, 3, 2], "cpu") | |
| def test_packed_to_padded_flat_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_packed_to_padded_helper([], device) | |
| def test_packed_to_padded_D1_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_packed_to_padded_helper([1], device) | |
| def test_packed_to_padded_D16_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_packed_to_padded_helper([16], device) | |
| def test_packed_to_padded_D16_9_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_packed_to_padded_helper([16, 9], device) | |
| def test_packed_to_padded_D16_3_2_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_packed_to_padded_helper([16, 3, 2], device) | |
| def _test_padded_to_packed_helper(self, dims, device): | |
| """ | |
| Check the results from packed_to_padded and PyTorch implementations | |
| are the same. | |
| """ | |
| meshes = self.init_meshes(16, 100, 300, device=device) | |
| mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() | |
| num_faces_per_mesh = meshes.num_faces_per_mesh() | |
| max_faces = num_faces_per_mesh.max().item() | |
| if len(dims) == 0: | |
| values = torch.rand((len(meshes), max_faces), device=device) | |
| else: | |
| values = torch.rand((len(meshes), max_faces, *dims), device=device) | |
| for i, num in enumerate(num_faces_per_mesh): | |
| values[i, num:] = 0 | |
| values.requires_grad = True | |
| values_torch = values.detach().clone() | |
| values_torch.requires_grad = True | |
| values_packed = padded_to_packed( | |
| values, mesh_to_faces_packed_first_idx, num_faces_per_mesh.sum().item() | |
| ) | |
| values_packed_torch = TestPackedToPadded.padded_to_packed_python( | |
| values_torch, | |
| mesh_to_faces_packed_first_idx, | |
| num_faces_per_mesh.sum().item(), | |
| device, | |
| ) | |
| # check forward | |
| self.assertClose(values_packed, values_packed_torch) | |
| if len(dims) > 0: | |
| values_packed_dim2 = padded_to_packed( | |
| values.transpose(1, 2), | |
| mesh_to_faces_packed_first_idx, | |
| num_faces_per_mesh.sum().item(), | |
| max_size_dim=2, | |
| ) | |
| # check forward | |
| self.assertClose(values_packed_dim2, values_packed_torch) | |
| # check backward | |
| if len(dims) == 0: | |
| grad_inputs = torch.rand((num_faces_per_mesh.sum().item()), device=device) | |
| else: | |
| grad_inputs = torch.rand( | |
| (num_faces_per_mesh.sum().item(), *dims), device=device | |
| ) | |
| values_packed.backward(grad_inputs) | |
| grad_outputs = values.grad | |
| values_packed_torch.backward(grad_inputs) | |
| grad_outputs_torch1 = values_torch.grad | |
| grad_outputs_torch2 = TestPackedToPadded.packed_to_padded_python( | |
| grad_inputs, mesh_to_faces_packed_first_idx, values.size(1), device=device | |
| ) | |
| self.assertClose(grad_outputs, grad_outputs_torch1) | |
| self.assertClose(grad_outputs, grad_outputs_torch2) | |
| def test_padded_to_packed_flat_cpu(self): | |
| self._test_padded_to_packed_helper([], "cpu") | |
| def test_padded_to_packed_D1_cpu(self): | |
| self._test_padded_to_packed_helper([1], "cpu") | |
| def test_padded_to_packed_D16_cpu(self): | |
| self._test_padded_to_packed_helper([16], "cpu") | |
| def test_padded_to_packed_D16_9_cpu(self): | |
| self._test_padded_to_packed_helper([16, 9], "cpu") | |
| def test_padded_to_packed_D16_3_2_cpu(self): | |
| self._test_padded_to_packed_helper([16, 3, 2], "cpu") | |
| def test_padded_to_packed_flat_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_padded_to_packed_helper([], device) | |
| def test_padded_to_packed_D1_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_padded_to_packed_helper([1], device) | |
| def test_padded_to_packed_D16_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_padded_to_packed_helper([16], device) | |
| def test_padded_to_packed_D16_9_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_padded_to_packed_helper([16, 9], device) | |
| def test_padded_to_packed_D16_3_2_cuda(self): | |
| device = get_random_cuda_device() | |
| self._test_padded_to_packed_helper([16, 3, 2], device) | |
| def packed_to_padded_with_init( | |
| num_meshes: int, num_verts: int, num_faces: int, num_d: int, device: str = "cpu" | |
| ): | |
| meshes = TestPackedToPadded.init_meshes( | |
| num_meshes, num_verts, num_faces, device | |
| ) | |
| faces = meshes.faces_packed() | |
| mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() | |
| max_faces = meshes.num_faces_per_mesh().max().item() | |
| if num_d == 0: | |
| values = torch.rand((faces.shape[0],), device=meshes.device) | |
| else: | |
| values = torch.rand((faces.shape[0], num_d), device=meshes.device) | |
| torch.cuda.synchronize() | |
| def out(): | |
| packed_to_padded(values, mesh_to_faces_packed_first_idx, max_faces) | |
| torch.cuda.synchronize() | |
| return out | |
| def packed_to_padded_with_init_torch( | |
| num_meshes: int, num_verts: int, num_faces: int, num_d: int, device: str = "cpu" | |
| ): | |
| meshes = TestPackedToPadded.init_meshes( | |
| num_meshes, num_verts, num_faces, device | |
| ) | |
| faces = meshes.faces_packed() | |
| mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() | |
| max_faces = meshes.num_faces_per_mesh().max().item() | |
| if num_d == 0: | |
| values = torch.rand((faces.shape[0],), device=meshes.device) | |
| else: | |
| values = torch.rand((faces.shape[0], num_d), device=meshes.device) | |
| torch.cuda.synchronize() | |
| def out(): | |
| TestPackedToPadded.packed_to_padded_python( | |
| values, mesh_to_faces_packed_first_idx, max_faces, device | |
| ) | |
| torch.cuda.synchronize() | |
| return out | |