detach for meshes, pointclouds, textures

Summary: Add `detach` for Meshes, Pointclouds, Textures

Reviewed By: nikhilaravi

Differential Revision: D23070418

fbshipit-source-id: 68671124ce114c4495d7ef3c944c9aac3d0db2d8

tests/test_meshes.py

@@ -20,6 +20,7 @@ class TestMeshes(TestCaseMixin, unittest.TestCase):
         max_f: int = 300,
         lists_to_tensors: bool = False,
         device: str = "cpu",
+        requires_grad: bool = False,
     ):
         """
         Function to generate a Meshes object of N meshes with
@@ -57,7 +58,12 @@ class TestMeshes(TestCaseMixin, unittest.TestCase):
 
         # Generate the actual vertices and faces.
         for i in range(num_meshes):
-            verts = torch.rand((v[i], 3), dtype=torch.float32, device=device)
+            verts = torch.rand(
+                (v[i], 3),
+                dtype=torch.float32,
+                device=device,
+                requires_grad=requires_grad,
+            )
             faces = torch.randint(
                 v[i], size=(f[i], 3), dtype=torch.int64, device=device
             )
@@ -353,6 +359,26 @@ class TestMeshes(TestCaseMixin, unittest.TestCase):
             self.assertSeparate(new_mesh.faces_padded(), mesh.faces_padded())
             self.assertSeparate(new_mesh.edges_packed(), mesh.edges_packed())
 
+    def test_detach(self):
+        N = 5
+        mesh = TestMeshes.init_mesh(N, 10, 100, requires_grad=True)
+        for force in [0, 1]:
+            if force:
+                # force mesh to have computed attributes
+                mesh.verts_packed()
+                mesh.edges_packed()
+                mesh.verts_padded()
+
+            new_mesh = mesh.detach()
+
+            self.assertFalse(new_mesh.verts_packed().requires_grad)
+            self.assertClose(new_mesh.verts_packed(), mesh.verts_packed())
+            self.assertTrue(new_mesh.verts_padded().requires_grad == False)
+            self.assertClose(new_mesh.verts_padded(), mesh.verts_padded())
+            for v, newv in zip(mesh.verts_list(), new_mesh.verts_list()):
+                self.assertTrue(newv.requires_grad == False)
+                self.assertClose(newv, v)
+
     def test_laplacian_packed(self):
         def naive_laplacian_packed(meshes):
             verts_packed = meshes.verts_packed()

tests/test_pointclouds.py

@@ -24,6 +24,7 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
         with_normals: bool = True,
         with_features: bool = True,
         min_points: int = 0,
+        requires_grad: bool = False,
     ):
         """
         Function to generate a Pointclouds object of N meshes with
@@ -49,16 +50,31 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
             p.fill_(p[0])
 
         points_list = [
-            torch.rand((i, 3), device=device, dtype=torch.float32) for i in p
+            torch.rand(
+                (i, 3), device=device, dtype=torch.float32, requires_grad=requires_grad
+            )
+            for i in p
         ]
         normals_list, features_list = None, None
         if with_normals:
             normals_list = [
-                torch.rand((i, 3), device=device, dtype=torch.float32) for i in p
+                torch.rand(
+                    (i, 3),
+                    device=device,
+                    dtype=torch.float32,
+                    requires_grad=requires_grad,
+                )
+                for i in p
             ]
         if with_features:
             features_list = [
-                torch.rand((i, channels), device=device, dtype=torch.float32) for i in p
+                torch.rand(
+                    (i, channels),
+                    device=device,
+                    dtype=torch.float32,
+                    requires_grad=requires_grad,
+                )
+                for i in p
             ]
 
         if lists_to_tensors:
@@ -382,6 +398,39 @@ class TestPointclouds(TestCaseMixin, unittest.TestCase):
 
             self.assertCloudsEqual(clouds, new_clouds)
 
+    def test_detach(self):
+        N = 5
+        for lists_to_tensors in (True, False):
+            clouds = self.init_cloud(
+                N, 100, 5, lists_to_tensors=lists_to_tensors, requires_grad=True
+            )
+            for force in (False, True):
+                if force:
+                    clouds.points_packed()
+
+                new_clouds = clouds.detach()
+
+                for cloud in new_clouds.points_list():
+                    self.assertTrue(cloud.requires_grad == False)
+                for normal in new_clouds.normals_list():
+                    self.assertTrue(normal.requires_grad == False)
+                for feats in new_clouds.features_list():
+                    self.assertTrue(feats.requires_grad == False)
+
+                for attrib in [
+                    "points_packed",
+                    "normals_packed",
+                    "features_packed",
+                    "points_padded",
+                    "normals_padded",
+                    "features_padded",
+                ]:
+                    self.assertTrue(
+                        getattr(new_clouds, attrib)().requires_grad == False
+                    )
+
+                self.assertCloudsEqual(clouds, new_clouds)
+
     def assertCloudsEqual(self, cloud1, cloud2):
         N = len(cloud1)
         self.assertEqual(N, len(cloud2))

tests/test_texturing.py

@@ -113,11 +113,37 @@ class TestTexturesVertex(TestCaseMixin, unittest.TestCase):
 
     def test_clone(self):
         tex = TexturesVertex(verts_features=torch.rand(size=(10, 100, 128)))
+        tex.verts_features_list()
         tex_cloned = tex.clone()
         self.assertSeparate(
             tex._verts_features_padded, tex_cloned._verts_features_padded
         )
+        self.assertClose(tex._verts_features_padded, tex_cloned._verts_features_padded)
         self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        for i in range(tex._N):
+            self.assertSeparate(
+                tex._verts_features_list[i], tex_cloned._verts_features_list[i]
+            )
+            self.assertClose(
+                tex._verts_features_list[i], tex_cloned._verts_features_list[i]
+            )
+
+    def test_detach(self):
+        tex = TexturesVertex(
+            verts_features=torch.rand(size=(10, 100, 128), requires_grad=True)
+        )
+        tex.verts_features_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._verts_features_padded.requires_grad)
+        self.assertClose(
+            tex_detached._verts_features_padded, tex._verts_features_padded
+        )
+        for i in range(tex._N):
+            self.assertClose(
+                tex._verts_features_list[i], tex_detached._verts_features_list[i]
+            )
+            self.assertFalse(tex_detached._verts_features_list[i].requires_grad)
 
     def test_extend(self):
         B = 10
@@ -278,9 +304,25 @@ class TestTexturesAtlas(TestCaseMixin, unittest.TestCase):
 
     def test_clone(self):
         tex = TexturesAtlas(atlas=torch.rand(size=(1, 10, 2, 2, 3)))
+        tex.atlas_list()
         tex_cloned = tex.clone()
         self.assertSeparate(tex._atlas_padded, tex_cloned._atlas_padded)
+        self.assertClose(tex._atlas_padded, tex_cloned._atlas_padded)
         self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        for i in range(tex._N):
+            self.assertSeparate(tex._atlas_list[i], tex_cloned._atlas_list[i])
+            self.assertClose(tex._atlas_list[i], tex_cloned._atlas_list[i])
+
+    def test_detach(self):
+        tex = TexturesAtlas(atlas=torch.rand(size=(1, 10, 2, 2, 3), requires_grad=True))
+        tex.atlas_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._atlas_padded.requires_grad)
+        self.assertClose(tex_detached._atlas_padded, tex._atlas_padded)
+        for i in range(tex._N):
+            self.assertFalse(tex_detached._atlas_list[i].requires_grad)
+            self.assertClose(tex._atlas_list[i], tex_detached._atlas_list[i])
 
     def test_extend(self):
         B = 10
@@ -478,11 +520,49 @@ class TestTexturesUV(TestCaseMixin, unittest.TestCase):
             faces_uvs=torch.rand(size=(5, 10, 3)),
             verts_uvs=torch.rand(size=(5, 15, 2)),
         )
+        tex.faces_uvs_list()
+        tex.verts_uvs_list()
         tex_cloned = tex.clone()
         self.assertSeparate(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
+        self.assertClose(tex._faces_uvs_padded, tex_cloned._faces_uvs_padded)
         self.assertSeparate(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
+        self.assertClose(tex._verts_uvs_padded, tex_cloned._verts_uvs_padded)
         self.assertSeparate(tex._maps_padded, tex_cloned._maps_padded)
+        self.assertClose(tex._maps_padded, tex_cloned._maps_padded)
         self.assertSeparate(tex.valid, tex_cloned.valid)
+        self.assertTrue(tex.valid.eq(tex_cloned.valid).all())
+        for i in range(tex._N):
+            self.assertSeparate(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
+            self.assertClose(tex._faces_uvs_list[i], tex_cloned._faces_uvs_list[i])
+            self.assertSeparate(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
+            self.assertClose(tex._verts_uvs_list[i], tex_cloned._verts_uvs_list[i])
+            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
+            self.assertSeparate(tex.maps_list()[i], tex_cloned.maps_list()[i])
+            self.assertClose(tex.maps_list()[i], tex_cloned.maps_list()[i])
+
+    def test_detach(self):
+        tex = TexturesUV(
+            maps=torch.ones((5, 16, 16, 3), requires_grad=True),
+            faces_uvs=torch.rand(size=(5, 10, 3)),
+            verts_uvs=torch.rand(size=(5, 15, 2)),
+        )
+        tex.faces_uvs_list()
+        tex.verts_uvs_list()
+        tex_detached = tex.detach()
+        self.assertFalse(tex_detached._maps_padded.requires_grad)
+        self.assertClose(tex._maps_padded, tex_detached._maps_padded)
+        self.assertFalse(tex_detached._verts_uvs_padded.requires_grad)
+        self.assertClose(tex._verts_uvs_padded, tex_detached._verts_uvs_padded)
+        self.assertFalse(tex_detached._faces_uvs_padded.requires_grad)
+        self.assertClose(tex._faces_uvs_padded, tex_detached._faces_uvs_padded)
+        for i in range(tex._N):
+            self.assertFalse(tex_detached._verts_uvs_list[i].requires_grad)
+            self.assertClose(tex._verts_uvs_list[i], tex_detached._verts_uvs_list[i])
+            self.assertFalse(tex_detached._faces_uvs_list[i].requires_grad)
+            self.assertClose(tex._faces_uvs_list[i], tex_detached._faces_uvs_list[i])
+            # tex._maps_list is not use anywhere so it's not stored. We call it explicitly
+            self.assertFalse(tex_detached.maps_list()[i].requires_grad)
+            self.assertClose(tex.maps_list()[i], tex_detached.maps_list()[i])
 
     def test_extend(self):
         B = 5