rasterizer.to without cameras

Summary: As reported in https://github.com/facebookresearch/pytorch3d/pull/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

pytorch3d/renderer/mesh/rasterizer.py

@@ -110,6 +110,7 @@ class MeshRasterizer(nn.Module):
 
     def to(self, device):
         # Manually move to device cameras as it is not a subclass of nn.Module
+        if self.cameras is not None:
             self.cameras = self.cameras.to(device)
         return self
 

pytorch3d/renderer/points/rasterizer.py

@@ -115,6 +115,7 @@ class PointsRasterizer(nn.Module):
 
     def to(self, device):
         # Manually move to device cameras as it is not a subclass of nn.Module
+        if self.cameras is not None:
             self.cameras = self.cameras.to(device)
         return self
 

tests/test_rasterizer.py

@@ -134,6 +134,12 @@ class TestMeshRasterizer(unittest.TestCase):
 
         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 @@ class TestPointRasterizer(unittest.TestCase):
             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)