423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-08-02 03:42:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

packed_to_padded now accepts all sizes

Browse Source 
Summary:
We need to make packing/unpacking in 2 places for mixed frame raysampling (metrics and raysampler) but those tensors that need to be unpacked/packed have more than two dimensions.
I could have reshaped and stored dimensions but this seems to just complicate code there with something which packed_to_padded should support.
I could have made a separate function for implicitron but it would confusing to have two different padded_to_packed functions inside pytorch3d codebase one of which does packing for (b, max) and (b, max, f) and the other for (b, max, …)

Reviewed By: bottler

Differential Revision: D39729026

fbshipit-source-id: 2bdebf290dcc6c316b7fe1aeee49bbb5255e508c
This commit is contained in:
Darijan Gudelj 2022-09-22 11:27:43 -07:00 committed by Facebook GitHub Bot
parent c2d876c9e8
commit f34da3d3b6
2 changed files with 86 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
pytorch3d/ops/packed_to_padded.py
View File

@ -65,7 +65,7 @@ def packed_to_padded(inputs, first_idxs, max_size):
Torch wrapper that handles allowed input shapes. See description below. Torch wrapper that handles allowed input shapes. See description below.
Args: Args:
inputs: FloatTensor of shape (F,) or (F, D), representing the packed inputs: FloatTensor of shape (F,) or (F, ...), representing the packed
batch tensor, e.g. areas for faces in a batch of meshes. batch tensor, e.g. areas for faces in a batch of meshes.
first_idxs: LongTensor of shape (N,) where N is the number of first_idxs: LongTensor of shape (N,) where N is the number of
elements in the batch and `first_idxs[i] = f` elements in the batch and `first_idxs[i] = f`
@ -73,7 +73,7 @@ def packed_to_padded(inputs, first_idxs, max_size):
max_size: Max length of an element in the batch. max_size: Max length of an element in the batch.
Returns: Returns:
inputs_padded: FloatTensor of shape (N, max_size) or (N, max_size, D) inputs_padded: FloatTensor of shape (N, max_size) or (N, max_size, ...)
where max_size is max of `sizes`. The values for batch element i where max_size is max of `sizes`. The values for batch element i
which start at `inputs[first_idxs[i]]` will be copied to which start at `inputs[first_idxs[i]]` will be copied to
`inputs_padded[i, :]`, with zeros padding out the extra inputs. `inputs_padded[i, :]`, with zeros padding out the extra inputs.
@ -83,15 +83,20 @@ def packed_to_padded(inputs, first_idxs, max_size):
(N, max_size, 1). (N, max_size, 1).
""" """
# if inputs is of shape (F,), reshape into (F, 1) # if inputs is of shape (F,), reshape into (F, 1)
flat = False input_shape = inputs.shape
if inputs.dim() == 1: n_dims = inputs.dim()
flat = True if n_dims == 1:
inputs = inputs.unsqueeze(1) inputs = inputs.unsqueeze(1)
else:
inputs = inputs.reshape(input_shape[0], -1)
inputs_padded = _PackedToPadded.apply(inputs, first_idxs, max_size) inputs_padded = _PackedToPadded.apply(inputs, first_idxs, max_size)
# if flat is True, reshape output to (N, max_size) from (N, max_size, 1) # if flat is True, reshape output to (N, max_size) from (N, max_size, 1)
if flat: # else reshape output to (N, max_size, ...)
inputs_padded = inputs_padded.squeeze(2) if n_dims == 1:
return inputs_padded return inputs_padded.squeeze(2)
if n_dims == 2:
return inputs_padded
return inputs_padded.view(*inputs_padded.shape[:2], *input_shape[1:])
class _PaddedToPacked(Function): class _PaddedToPacked(Function):
@ -147,7 +152,7 @@ def padded_to_packed(inputs, first_idxs, num_inputs):
Torch wrapper that handles allowed input shapes. See description below. Torch wrapper that handles allowed input shapes. See description below.
Args: Args:
inputs: FloatTensor of shape (N, max_size) or (N, max_size, D), inputs: FloatTensor of shape (N, max_size) or (N, max_size, ...),
representing the padded tensor, e.g. areas for faces in a batch of representing the padded tensor, e.g. areas for faces in a batch of
meshes. meshes.
first_idxs: LongTensor of shape (N,) where N is the number of first_idxs: LongTensor of shape (N,) where N is the number of
@ -156,20 +161,25 @@ def padded_to_packed(inputs, first_idxs, num_inputs):
num_inputs: Number of packed entries (= F) num_inputs: Number of packed entries (= F)
Returns: Returns:
inputs_packed: FloatTensor of shape (F,) or (F, D) where inputs_packed: FloatTensor of shape (F,) or (F, ...) where
`inputs_packed[first_idx[i]:] = inputs[i, :]`. `inputs_packed[first_idx[i]:first_idx[i+1]] = inputs[i, :]`.
To handle the allowed input shapes, we convert the inputs tensor of shape To handle the allowed input shapes, we convert the inputs tensor of shape
(N, max_size) to (N, max_size, 1). We reshape the output back to (F,) from (N, max_size) to (N, max_size, 1). We reshape the output back to (F,) from
(F, 1). (F, 1).
""" """
# if inputs is of shape (N, max_size), reshape into (N, max_size, 1)) # if inputs is of shape (N, max_size), reshape into (N, max_size, 1))
flat = False input_shape = inputs.shape
if inputs.dim() == 2: n_dims = inputs.dim()
flat = True if n_dims == 2:
inputs = inputs.unsqueeze(2) inputs = inputs.unsqueeze(2)
else:
inputs = inputs.reshape(*input_shape[:2], -1)
inputs_packed = _PaddedToPacked.apply(inputs, first_idxs, num_inputs) inputs_packed = _PaddedToPacked.apply(inputs, first_idxs, num_inputs)
# if flat is True, reshape output to (F,) from (F, 1) # if input is flat, reshape output to (F,) from (F, 1)
if flat: # else reshape output to (F, ...)
inputs_packed = inputs_packed.squeeze(1) if n_dims == 2:
return inputs_packed return inputs_packed.squeeze(1)
if n_dims == 3:
return inputs_packed
return inputs_packed.view(-1, *input_shape[2:])

102
tests/test_packed_to_padded.py
View File

@ -45,18 +45,19 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
PyTorch implementation of packed_to_padded function. PyTorch implementation of packed_to_padded function.
""" """
num_meshes = first_idxs.size(0) num_meshes = first_idxs.size(0)
D = inputs.shape[1] if inputs.dim() == 2 else 0 if inputs.dim() == 1:
if D == 0:
inputs_padded = torch.zeros((num_meshes, max_size), device=device) inputs_padded = torch.zeros((num_meshes, max_size), device=device)
else: else:
inputs_padded = torch.zeros((num_meshes, max_size, D), device=device) inputs_padded = torch.zeros(
(num_meshes, max_size, *inputs.shape[1:]), device=device
)
for m in range(num_meshes): for m in range(num_meshes):
s = first_idxs[m] s = first_idxs[m]
if m == num_meshes - 1: if m == num_meshes - 1:
f = inputs.shape[0] f = inputs.shape[0]
else: else:
f = first_idxs[m + 1] f = first_idxs[m + 1]
inputs_padded[m, :f] = inputs[s:f] inputs_padded[m, : f - s] = inputs[s:f]
return inputs_padded return inputs_padded
@ -66,22 +67,21 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
PyTorch implementation of padded_to_packed function. PyTorch implementation of padded_to_packed function.
""" """
num_meshes = inputs.size(0) num_meshes = inputs.size(0)
D = inputs.shape[2] if inputs.dim() == 3 else 0 if inputs.dim() == 2:
if D == 0:
inputs_packed = torch.zeros((num_inputs,), device=device) inputs_packed = torch.zeros((num_inputs,), device=device)
else: else:
inputs_packed = torch.zeros((num_inputs, D), device=device) inputs_packed = torch.zeros((num_inputs, *inputs.shape[2:]), device=device)
for m in range(num_meshes): for m in range(num_meshes):
s = first_idxs[m] s = first_idxs[m]
if m == num_meshes - 1: if m == num_meshes - 1:
f = num_inputs f = num_inputs
else: else:
f = first_idxs[m + 1] f = first_idxs[m + 1]
inputs_packed[s:f] = inputs[m, :f] inputs_packed[s:f] = inputs[m, : f - s]
return inputs_packed return inputs_packed
def _test_packed_to_padded_helper(self, D, device): def _test_packed_to_padded_helper(self, dims, device):
""" """
Check the results from packed_to_padded and PyTorch implementations Check the results from packed_to_padded and PyTorch implementations
are the same. are the same.
@ -91,10 +91,12 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
max_faces = meshes.num_faces_per_mesh().max().item() max_faces = meshes.num_faces_per_mesh().max().item()
if D == 0: if len(dims) == 0:
values = torch.rand((faces.shape[0],), device=device, requires_grad=True) values = torch.rand((faces.shape[0],), device=device, requires_grad=True)
else: else:
values = torch.rand((faces.shape[0], D), device=device, requires_grad=True) values = torch.rand(
(faces.shape[0], *dims), device=device, requires_grad=True
)
values_torch = values.detach().clone() values_torch = values.detach().clone()
values_torch.requires_grad = True values_torch.requires_grad = True
values_padded = packed_to_padded( values_padded = packed_to_padded(
@ -107,10 +109,10 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
self.assertClose(values_padded, values_padded_torch) self.assertClose(values_padded, values_padded_torch)
# check backward # check backward
if D == 0: if len(dims) == 0:
grad_inputs = torch.rand((len(meshes), max_faces), device=device) grad_inputs = torch.rand((len(meshes), max_faces), device=device)
else: else:
grad_inputs = torch.rand((len(meshes), max_faces, D), device=device) grad_inputs = torch.rand((len(meshes), max_faces, *dims), device=device)
values_padded.backward(grad_inputs) values_padded.backward(grad_inputs)
grad_outputs = values.grad grad_outputs = values.grad
values_padded_torch.backward(grad_inputs) values_padded_torch.backward(grad_inputs)
@ -122,27 +124,41 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
self.assertClose(grad_outputs, grad_outputs_torch2) self.assertClose(grad_outputs, grad_outputs_torch2)
def test_packed_to_padded_flat_cpu(self): def test_packed_to_padded_flat_cpu(self):
self._test_packed_to_padded_helper(0, "cpu") self._test_packed_to_padded_helper([], "cpu")
def test_packed_to_padded_D1_cpu(self): def test_packed_to_padded_D1_cpu(self):
self._test_packed_to_padded_helper(1, "cpu") self._test_packed_to_padded_helper([1], "cpu")
def test_packed_to_padded_D16_cpu(self): def test_packed_to_padded_D16_cpu(self):
self._test_packed_to_padded_helper(16, "cpu") 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): def test_packed_to_padded_flat_cuda(self):
device = get_random_cuda_device() device = get_random_cuda_device()
self._test_packed_to_padded_helper(0, device) self._test_packed_to_padded_helper([], device)
def test_packed_to_padded_D1_cuda(self): def test_packed_to_padded_D1_cuda(self):
device = get_random_cuda_device() device = get_random_cuda_device()
self._test_packed_to_padded_helper(1, device) self._test_packed_to_padded_helper([1], device)
def test_packed_to_padded_D16_cuda(self): def test_packed_to_padded_D16_cuda(self):
device = get_random_cuda_device() device = get_random_cuda_device()
self._test_packed_to_padded_helper(16, device) self._test_packed_to_padded_helper([16], device)
def _test_padded_to_packed_helper(self, D, 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 Check the results from packed_to_padded and PyTorch implementations
are the same. are the same.
@ -151,10 +167,10 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx() mesh_to_faces_packed_first_idx = meshes.mesh_to_faces_packed_first_idx()
num_faces_per_mesh = meshes.num_faces_per_mesh() num_faces_per_mesh = meshes.num_faces_per_mesh()
max_faces = num_faces_per_mesh.max().item() max_faces = num_faces_per_mesh.max().item()
if D == 0: if len(dims) == 0:
values = torch.rand((len(meshes), max_faces), device=device) values = torch.rand((len(meshes), max_faces), device=device)
else: else:
values = torch.rand((len(meshes), max_faces, D), device=device) values = torch.rand((len(meshes), max_faces, *dims), device=device)
for i, num in enumerate(num_faces_per_mesh): for i, num in enumerate(num_faces_per_mesh):
values[i, num:] = 0 values[i, num:] = 0
values.requires_grad = True values.requires_grad = True
@ -173,11 +189,11 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
self.assertClose(values_packed, values_packed_torch) self.assertClose(values_packed, values_packed_torch)
# check backward # check backward
if D == 0: if len(dims) == 0:
grad_inputs = torch.rand((num_faces_per_mesh.sum().item()), device=device) grad_inputs = torch.rand((num_faces_per_mesh.sum().item()), device=device)
else: else:
grad_inputs = torch.rand( grad_inputs = torch.rand(
(num_faces_per_mesh.sum().item(), D), device=device (num_faces_per_mesh.sum().item(), *dims), device=device
) )
values_packed.backward(grad_inputs) values_packed.backward(grad_inputs)
grad_outputs = values.grad grad_outputs = values.grad
@ -190,41 +206,39 @@ class TestPackedToPadded(TestCaseMixin, unittest.TestCase):
self.assertClose(grad_outputs, grad_outputs_torch2) self.assertClose(grad_outputs, grad_outputs_torch2)
def test_padded_to_packed_flat_cpu(self): def test_padded_to_packed_flat_cpu(self):
self._test_padded_to_packed_helper(0, "cpu") self._test_padded_to_packed_helper([], "cpu")
def test_padded_to_packed_D1_cpu(self): def test_padded_to_packed_D1_cpu(self):
self._test_padded_to_packed_helper(1, "cpu") self._test_padded_to_packed_helper([1], "cpu")
def test_padded_to_packed_D16_cpu(self): def test_padded_to_packed_D16_cpu(self):
self._test_padded_to_packed_helper(16, "cpu") 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): def test_padded_to_packed_flat_cuda(self):
device = get_random_cuda_device() device = get_random_cuda_device()
self._test_padded_to_packed_helper(0, device) self._test_padded_to_packed_helper([], device)
def test_padded_to_packed_D1_cuda(self): def test_padded_to_packed_D1_cuda(self):
device = get_random_cuda_device() device = get_random_cuda_device()
self._test_padded_to_packed_helper(1, device) self._test_padded_to_packed_helper([1], device)
def test_padded_to_packed_D16_cuda(self): def test_padded_to_packed_D16_cuda(self):
device = get_random_cuda_device() device = get_random_cuda_device()
self._test_padded_to_packed_helper(16, device) self._test_padded_to_packed_helper([16], device)
def test_invalid_inputs_shapes(self, device="cuda:0"): def test_padded_to_packed_D16_9_cuda(self):
with self.assertRaisesRegex(ValueError, "input can only be 2-dimensional."): device = get_random_cuda_device()
values = torch.rand((100, 50, 2), device=device) self._test_padded_to_packed_helper([16, 9], device)
first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
packed_to_padded(values, first_idxs, 100)
with self.assertRaisesRegex(ValueError, "input can only be 3-dimensional."): def test_padded_to_packed_D16_3_2_cuda(self):
values = torch.rand((100,), device=device) device = get_random_cuda_device()
first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device) self._test_padded_to_packed_helper([16, 3, 2], device)
padded_to_packed(values, first_idxs, 20)
with self.assertRaisesRegex(ValueError, "input can only be 3-dimensional."):
values = torch.rand((100, 50, 2, 2), device=device)
first_idxs = torch.tensor([0, 80], dtype=torch.int64, device=device)
padded_to_packed(values, first_idxs, 20)
@staticmethod @staticmethod
def packed_to_padded_with_init( def packed_to_padded_with_init(