pytorch3d icon indicating copy to clipboard operation
pytorch3d copied to clipboard

Meshes.update_padded_

Open timlod opened this issue 2 years ago • 2 comments

🚀 Feature

Currently, update_padded only exists as an out-of-place operation which shallow copies attributes onto a new object. It would be nice to have an in-place version that just modifies the vertices of the existing Meshes object.

Motivation

I'm updating a mesh, but I'm not generating the new vertices through pytorch3d, or some method such that I could use offset_verts_, but an auxiliary model. I would like to reuse the existing Meshes object for better efficiency and to allow the use of side effects when convenient.

Pitch

I may be missing some things (like updating normals if they'd been computed), but this could be enough:

    def update_padded_(self, new_verts_padded):
        """
        In-place version of update_padded.

        Args:
            new_verts_padded: FloatTensor of shape (N, V, 3)

        Returns:
            Meshes with updated padded representations
        """

        def check_shapes(x, size):
            if x.shape[0] != size[0]:
                raise ValueError("new values must have the same batch dimension.")
            if x.shape[1] != size[1]:
                raise ValueError("new values must have the same number of points.")
            if x.shape[2] != size[2]:
                raise ValueError("new values must have the same dimension.")

        self._verts_padded = new_verts_padded

        # update verts/faces packed if they are computed in self
        if self._verts_packed is not None:
            # update verts_packed
            pad_to_packed = self.verts_padded_to_packed_idx()
            new_verts_packed = new_verts_padded.reshape(-1, 3)[pad_to_packed, :]
            self._verts_packed = new_verts_packed
            self._verts_padded_to_packed_idx = pad_to_packed

        return self

There's an obvious huge speed-up involved when not creating a new object.

timlod avatar Dec 09 '21 09:12 timlod

Have you measured the speed up?

In the applications we care about, saving GPU compute and memory is usually more important than saving Python objects. Most of the time the GPU will be running asyncronous operations which form the critical path. I think your change typically won't change what the GPU does.

I think allowing people to modify the object in place makes it easier to silently do the wrong thing with autograd. I haven't got a simple example. If verts are changing during training, we usually want to encourage the user to recreate the Meshes objects in each optimization step. offset_verts_ and scale_verts_ are the only current exceptions.

bottler avatar Dec 09 '21 10:12 bottler

I only measured the two functions against each other, not in a rendering pipeline. I guess my last sentence is an overstatement, given that, as you say, async ops on the GPU may take place while this operation is happening. I just measured the speed-up when rasterizing after the operation, and it's now minimal (~5k verts, half a ms speed-up).

My main reason for wanting this in any case is the in-place modification to be able to use side-effects if wanted, but I understand if you don't want to allow this in general (given the pitfalls it can create).

timlod avatar Dec 09 '21 10:12 timlod