rasterizer.to without cameras

bottler · facebook-github-bot · commit c371a9a6ccec · 2022-03-08T23:54:38.000-08:00
Summary: As reported in #1100, a rasterizer couldn't be moved if it was missing the optional cameras member. Fix that. This matters because the renderer.to calls rasterizer.to, so this to() could be called even by a user who never sets a cameras member. Reviewed By: nikhilaravi Differential Revision: D34643841 fbshipit-source-id: 7e26e32e8bc585eb1ee533052754a7b59bc7467a
diff --git a/pytorch3d/renderer/mesh/rasterizer.py b/pytorch3d/renderer/mesh/rasterizer.py
@@ -110,7 +110,8 @@ def __init__(self, cameras=None, raster_settings=None) -> None:
 
     def to(self, device):
         # Manually move to device cameras as it is not a subclass of nn.Module
-        self.cameras = self.cameras.to(device)
+        if self.cameras is not None:
+            self.cameras = self.cameras.to(device)
         return self
 
     def transform(self, meshes_world, **kwargs) -> torch.Tensor:
diff --git a/pytorch3d/renderer/points/rasterizer.py b/pytorch3d/renderer/points/rasterizer.py
@@ -115,7 +115,8 @@ def transform(self, point_clouds, **kwargs) -> torch.Tensor:
 
     def to(self, device):
         # Manually move to device cameras as it is not a subclass of nn.Module
-        self.cameras = self.cameras.to(device)
+        if self.cameras is not None:
+            self.cameras = self.cameras.to(device)
         return self
 
     def forward(self, point_clouds, **kwargs) -> PointFragments:
diff --git a/tests/test_rasterizer.py b/tests/test_rasterizer.py
@@ -134,6 +134,12 @@ def test_simple_sphere(self):
 
         self.assertTrue(torch.allclose(image, image_ref))
 
+    def test_simple_to(self):
+        # Check that to() works without a cameras object.
+        device = torch.device("cuda:0")
+        rasterizer = MeshRasterizer()
+        rasterizer.to(device)
+
 
 class TestPointRasterizer(unittest.TestCase):
     def test_simple_sphere(self):
@@ -203,3 +209,9 @@ def test_simple_sphere(self):
             image[image >= 0] = 1.0
             image[image < 0] = 0.0
             self.assertTrue(torch.allclose(image, image_ref[..., 0]))
+
+    def test_simple_to(self):
+        # Check that to() works without a cameras object.
+        device = torch.device("cuda:0")
+        rasterizer = PointsRasterizer()
+        rasterizer.to(device)
