Add MeshRasterizerOpenGL

Summary:
Adding MeshRasterizerOpenGL, a faster alternative to MeshRasterizer. The new rasterizer follows the ideas from "Differentiable Surface Rendering via non-Differentiable Sampling".

The new rasterizer 20x faster on a 2M face mesh (try pose optimization on Nefertiti from https://www.cs.cmu.edu/~kmcrane/Projects/ModelRepository/!). The larger the mesh, the larger the speedup.

There are two main disadvantages:
* The new rasterizer works with an OpenGL backend, so requires pycuda.gl and pyopengl installed (though we avoided writing any C++ code, everything is in Python!)
* The new rasterizer is non-differentiable. However, you can still differentiate the rendering function if you use if with the new SplatterPhongShader which we recently added to PyTorch3D (see the original paper cited above).

Reviewed By: patricklabatut, jcjohnson

Differential Revision: D37698816

fbshipit-source-id: 54d120639d3cb001f096237807e54aced0acda25

tests/implicitron/test_build.py

@@ -26,8 +26,15 @@ class TestBuild(unittest.TestCase):
                     sys.modules.pop(module, None)
 
             root_dir = get_pytorch3d_dir() / "pytorch3d"
+            # Exclude opengl-related files, as Implicitron is decoupled from opengl
+            # components which will not work without adding a dep on pytorch3d_opengl.
             for module_file in root_dir.glob("**/*.py"):
-                if module_file.stem in ("__init__", "plotly_vis", "opengl_utils"):
+                if module_file.stem in (
+                    "__init__",
+                    "plotly_vis",
+                    "opengl_utils",
+                    "rasterizer_opengl",
+                ):
                     continue
                 relative_module = str(module_file.relative_to(root_dir))[:-3]
                 module = "pytorch3d." + relative_module.replace("/", ".")

tests/test_rasterize_rectangle_images.py

@@ -11,15 +11,18 @@ import numpy as np
 import torch
 from PIL import Image
 from pytorch3d.io import load_obj
-from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
-from pytorch3d.renderer.lighting import PointLights
-from pytorch3d.renderer.materials import Materials
-from pytorch3d.renderer.mesh import (
+from pytorch3d.renderer import (
     BlendParams,
+    FoVPerspectiveCameras,
+    look_at_view_transform,
+    Materials,
     MeshRasterizer,
+    MeshRasterizerOpenGL,
     MeshRenderer,
+    PointLights,
     RasterizationSettings,
     SoftPhongShader,
+    SplatterPhongShader,
     TexturesUV,
 )
 from pytorch3d.renderer.mesh.rasterize_meshes import (
@@ -454,6 +457,12 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase):
             )
 
     def test_render_cow(self):
+        self._render_cow(MeshRasterizer)
+
+    def test_render_cow_opengl(self):
+        self._render_cow(MeshRasterizerOpenGL)
+
+    def _render_cow(self, rasterizer_type):
         """
         Test a larger textured mesh is rendered correctly in a non square image.
         """
@@ -473,38 +482,55 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase):
         mesh = Meshes(verts=[verts], faces=[faces.verts_idx], textures=textures)
 
         # Init rasterizer settings
-        R, T = look_at_view_transform(2.7, 0, 180)
+        R, T = look_at_view_transform(1.2, 0, 90)
         cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
 
         raster_settings = RasterizationSettings(
-            image_size=(512, 1024), blur_radius=0.0, faces_per_pixel=1
+            image_size=(500, 800), blur_radius=0.0, faces_per_pixel=1
         )
 
         # Init shader settings
         materials = Materials(device=device)
         lights = PointLights(device=device)
         lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None]
-        blend_params = BlendParams(
-            sigma=1e-1,
-            gamma=1e-4,
-            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
-        )
 
         # Init renderer
-        renderer = MeshRenderer(
-            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
-            shader=SoftPhongShader(
+        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
+        if rasterizer_type == MeshRasterizer:
+            blend_params = BlendParams(
+                sigma=1e-1,
+                gamma=1e-4,
+                background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
+            )
+            shader = SoftPhongShader(
                 lights=lights,
                 cameras=cameras,
                 materials=materials,
                 blend_params=blend_params,
-            ),
-        )
+            )
+        else:
+            blend_params = BlendParams(
+                sigma=0.5,
+                background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
+            )
+            shader = SplatterPhongShader(
+                lights=lights,
+                cameras=cameras,
+                materials=materials,
+                blend_params=blend_params,
+            )
+
+        renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
 
         # Load reference image
-        image_ref = load_rgb_image("test_cow_image_rectangle.png", DATA_DIR)
+        image_ref = load_rgb_image(
+            f"test_cow_image_rectangle_{rasterizer_type.__name__}.png", DATA_DIR
+        )
 
         for bin_size in [0, None]:
+            if bin_size == 0 and rasterizer_type == MeshRasterizerOpenGL:
+                continue
+
             # Check both naive and coarse to fine produce the same output.
             renderer.rasterizer.raster_settings.bin_size = bin_size
             images = renderer(mesh)
@@ -512,7 +538,8 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase):
 
             if DEBUG:
                 Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
-                    DATA_DIR / "DEBUG_cow_image_rectangle.png"
+                    DATA_DIR
+                    / f"DEBUG_cow_image_rectangle_{rasterizer_type.__name__}.png"
                 )
 
             # NOTE some pixels can be flaky

tests/test_rasterizer.py

@@ -10,16 +10,29 @@ import unittest
 import numpy as np
 import torch
 from PIL import Image
-from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
-from pytorch3d.renderer.mesh.rasterizer import MeshRasterizer, RasterizationSettings
-from pytorch3d.renderer.points.rasterizer import (
+from pytorch3d.renderer import (
+    FoVOrthographicCameras,
+    FoVPerspectiveCameras,
+    look_at_view_transform,
+    MeshRasterizer,
+    MeshRasterizerOpenGL,
+    OrthographicCameras,
+    PerspectiveCameras,
     PointsRasterizationSettings,
     PointsRasterizer,
+    RasterizationSettings,
+)
+from pytorch3d.renderer.opengl.rasterizer_opengl import (
+    _check_cameras,
+    _check_raster_settings,
+    _convert_meshes_to_gl_ndc,
+    _parse_and_verify_image_size,
 )
 from pytorch3d.structures import Pointclouds
+from pytorch3d.structures.meshes import Meshes
 from pytorch3d.utils.ico_sphere import ico_sphere
 
-from .common_testing import get_tests_dir
+from .common_testing import get_tests_dir, TestCaseMixin
 
 
 DATA_DIR = get_tests_dir() / "data"
@@ -36,8 +49,14 @@ def convert_image_to_binary_mask(filename):
 
 class TestMeshRasterizer(unittest.TestCase):
     def test_simple_sphere(self):
+        self._simple_sphere(MeshRasterizer)
+
+    def test_simple_sphere_opengl(self):
+        self._simple_sphere(MeshRasterizerOpenGL)
+
+    def _simple_sphere(self, rasterizer_type):
         device = torch.device("cuda:0")
-        ref_filename = "test_rasterized_sphere.png"
+        ref_filename = f"test_rasterized_sphere_{rasterizer_type.__name__}.png"
         image_ref_filename = DATA_DIR / ref_filename
 
         # Rescale image_ref to the 0 - 1 range and convert to a binary mask.
@@ -54,7 +73,7 @@ class TestMeshRasterizer(unittest.TestCase):
         )
 
         # Init rasterizer
-        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
+        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
 
         ####################################
         # 1. Test rasterizing a single mesh
@@ -68,7 +87,8 @@ class TestMeshRasterizer(unittest.TestCase):
 
         if DEBUG:
             Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
-                DATA_DIR / "DEBUG_test_rasterized_sphere.png"
+                DATA_DIR
+                / f"DEBUG_test_rasterized_sphere_{rasterizer_type.__name__}.png"
             )
 
         self.assertTrue(torch.allclose(image, image_ref))
@@ -90,20 +110,21 @@ class TestMeshRasterizer(unittest.TestCase):
         #  3. Test that passing kwargs to rasterizer works.
         ####################################################
 
-        #  Change the view transform to zoom in.
-        R, T = look_at_view_transform(2.0, 0, 0, device=device)
+        #  Change the view transform to zoom out.
+        R, T = look_at_view_transform(20.0, 0, 0, device=device)
         fragments = rasterizer(sphere_mesh, R=R, T=T)
         image = fragments.pix_to_face[0, ..., 0].squeeze().cpu()
         image[image >= 0] = 1.0
         image[image < 0] = 0.0
 
-        ref_filename = "test_rasterized_sphere_zoom.png"
+        ref_filename = f"test_rasterized_sphere_zoom_{rasterizer_type.__name__}.png"
         image_ref_filename = DATA_DIR / ref_filename
         image_ref = convert_image_to_binary_mask(image_ref_filename)
 
         if DEBUG:
             Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
-                DATA_DIR / "DEBUG_test_rasterized_sphere_zoom.png"
+                DATA_DIR
+                / f"DEBUG_test_rasterized_sphere_zoom_{rasterizer_type.__name__}.png"
             )
         self.assertTrue(torch.allclose(image, image_ref))
 
@@ -112,7 +133,7 @@ class TestMeshRasterizer(unittest.TestCase):
         ##################################
 
         # Create a new empty rasterizer:
-        rasterizer = MeshRasterizer()
+        rasterizer = rasterizer_type(raster_settings=raster_settings)
 
         # Check that omitting the cameras in both initialization
         # and the forward pass throws an error:
@@ -120,9 +141,7 @@ class TestMeshRasterizer(unittest.TestCase):
             rasterizer(sphere_mesh)
 
         # Now pass in the cameras as a kwarg
-        fragments = rasterizer(
-            sphere_mesh, cameras=cameras, raster_settings=raster_settings
-        )
+        fragments = rasterizer(sphere_mesh, cameras=cameras)
         image = fragments.pix_to_face[0, ..., 0].squeeze().cpu()
         # Convert pix_to_face to a binary mask
         image[image >= 0] = 1.0
@@ -130,7 +149,8 @@ class TestMeshRasterizer(unittest.TestCase):
 
         if DEBUG:
             Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
-                DATA_DIR / "DEBUG_test_rasterized_sphere.png"
+                DATA_DIR
+                / f"DEBUG_test_rasterized_sphere_{rasterizer_type.__name__}.png"
             )
 
         self.assertTrue(torch.allclose(image, image_ref))
@@ -141,6 +161,187 @@ class TestMeshRasterizer(unittest.TestCase):
         rasterizer = MeshRasterizer()
         rasterizer.to(device)
 
+        rasterizer = MeshRasterizerOpenGL()
+        rasterizer.to(device)
+
+    def test_compare_rasterizers(self):
+        device = torch.device("cuda:0")
+
+        # Init rasterizer settings
+        R, T = look_at_view_transform(2.7, 0, 0)
+        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
+        raster_settings = RasterizationSettings(
+            image_size=512,
+            blur_radius=0.0,
+            faces_per_pixel=1,
+            bin_size=0,
+            perspective_correct=True,
+        )
+        from pytorch3d.io import load_obj
+        from pytorch3d.renderer import TexturesAtlas
+
+        from .common_testing import get_pytorch3d_dir
+
+        TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
+        obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
+
+        # Load mesh and texture as a per face texture atlas.
+        verts, faces, aux = load_obj(
+            obj_filename,
+            device=device,
+            load_textures=True,
+            create_texture_atlas=True,
+            texture_atlas_size=8,
+            texture_wrap=None,
+        )
+        atlas = aux.texture_atlas
+        mesh = Meshes(
+            verts=[verts],
+            faces=[faces.verts_idx],
+            textures=TexturesAtlas(atlas=[atlas]),
+        )
+
+        # Rasterize using both rasterizers and compare results.
+        rasterizer = MeshRasterizerOpenGL(
+            cameras=cameras, raster_settings=raster_settings
+        )
+        fragments_opengl = rasterizer(mesh)
+
+        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
+        fragments = rasterizer(mesh)
+
+        # Ensure that 99.9% of bary_coords is at most 0.001 different.
+        self.assertLess(
+            torch.quantile(
+                (fragments.bary_coords - fragments_opengl.bary_coords).abs(), 0.999
+            ),
+            0.001,
+        )
+
+        # Ensure that 99.9% of zbuf vals is at most 0.001 different.
+        self.assertLess(
+            torch.quantile((fragments.zbuf - fragments_opengl.zbuf).abs(), 0.999), 0.001
+        )
+
+        # Ensure that 99.99% of pix_to_face is identical.
+        self.assertEqual(
+            torch.quantile(
+                (fragments.pix_to_face != fragments_opengl.pix_to_face).float(), 0.9999
+            ),
+            0,
+        )
+
+
+class TestMeshRasterizerOpenGLUtils(TestCaseMixin, unittest.TestCase):
+    def setUp(self):
+        verts = torch.tensor(
+            [[-1, 1, 0], [1, 1, 0], [1, -1, 0]], dtype=torch.float32
+        ).cuda()
+        faces = torch.tensor([[0, 1, 2]]).cuda()
+        self.meshes_world = Meshes(verts=[verts], faces=[faces])
+
+    # Test various utils specific to the OpenGL rasterizer. Full "integration tests"
+    # live in test_render_meshes and test_render_multigpu.
+    def test_check_cameras(self):
+        _check_cameras(FoVPerspectiveCameras())
+        _check_cameras(FoVPerspectiveCameras())
+        with self.assertRaisesRegex(ValueError, "Cameras must be specified"):
+            _check_cameras(None)
+        with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
+            _check_cameras(PerspectiveCameras())
+        with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
+            _check_cameras(OrthographicCameras())
+
+        MeshRasterizerOpenGL(FoVPerspectiveCameras().cuda())(self.meshes_world)
+        MeshRasterizerOpenGL(FoVOrthographicCameras().cuda())(self.meshes_world)
+        MeshRasterizerOpenGL()(
+            self.meshes_world, cameras=FoVPerspectiveCameras().cuda()
+        )
+
+        with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
+            MeshRasterizerOpenGL(PerspectiveCameras().cuda())(self.meshes_world)
+        with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
+            MeshRasterizerOpenGL(OrthographicCameras().cuda())(self.meshes_world)
+        with self.assertRaisesRegex(ValueError, "Cameras must be specified"):
+            MeshRasterizerOpenGL()(self.meshes_world)
+
+    def test_check_raster_settings(self):
+        raster_settings = RasterizationSettings()
+        raster_settings.faces_per_pixel = 100
+        with self.assertWarnsRegex(UserWarning, ".* one face per pixel"):
+            _check_raster_settings(raster_settings)
+
+        with self.assertWarnsRegex(UserWarning, ".* one face per pixel"):
+            MeshRasterizerOpenGL(raster_settings=raster_settings)(
+                self.meshes_world, cameras=FoVPerspectiveCameras().cuda()
+            )
+
+    def test_convert_meshes_to_gl_ndc_square_img(self):
+        R, T = look_at_view_transform(1, 90, 180)
+        cameras = FoVOrthographicCameras(R=R, T=T).cuda()
+
+        meshes_gl_ndc = _convert_meshes_to_gl_ndc(
+            self.meshes_world, (100, 100), cameras
+        )
+
+        # After look_at_view_transform rotating 180 deg around z-axis, we recover
+        # the original coordinates. After additionally elevating the view by 90
+        # deg, we "zero out" the y-coordinate. Finally, we negate the x and y axes
+        # to adhere to OpenGL conventions (which go against the PyTorch3D convention).
+        self.assertClose(
+            meshes_gl_ndc.verts_list()[0],
+            torch.tensor(
+                [[-1, 0, 0], [1, 0, 0], [1, 0, 2]], dtype=torch.float32
+            ).cuda(),
+            atol=1e-5,
+        )
+
+    def test_parse_and_verify_image_size(self):
+        img_size = _parse_and_verify_image_size(512)
+        self.assertEqual(img_size, (512, 512))
+
+        img_size = _parse_and_verify_image_size((2047, 10))
+        self.assertEqual(img_size, (2047, 10))
+
+        img_size = _parse_and_verify_image_size((10, 2047))
+        self.assertEqual(img_size, (10, 2047))
+
+        with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
+            _parse_and_verify_image_size((2049, 512))
+
+        with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
+            _parse_and_verify_image_size((512, 2049))
+
+        with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
+            _parse_and_verify_image_size((2049, 2049))
+
+        rasterizer = MeshRasterizerOpenGL(FoVPerspectiveCameras().cuda())
+        raster_settings = RasterizationSettings()
+
+        raster_settings.image_size = 512
+        fragments = rasterizer(self.meshes_world, raster_settings=raster_settings)
+        self.assertEqual(fragments.pix_to_face.shape, torch.Size([1, 512, 512, 1]))
+
+        raster_settings.image_size = (2047, 10)
+        fragments = rasterizer(self.meshes_world, raster_settings=raster_settings)
+        self.assertEqual(fragments.pix_to_face.shape, torch.Size([1, 2047, 10, 1]))
+
+        raster_settings.image_size = (10, 2047)
+        fragments = rasterizer(self.meshes_world, raster_settings=raster_settings)
+        self.assertEqual(fragments.pix_to_face.shape, torch.Size([1, 10, 2047, 1]))
+
+        with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
+            raster_settings.image_size = (2049, 512)
+            rasterizer(self.meshes_world, raster_settings=raster_settings)
+
+        with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
+            raster_settings.image_size = (512, 2049)
+            rasterizer(self.meshes_world, raster_settings=raster_settings)
+
+        with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
+            raster_settings.image_size = (2049, 2049)
+            rasterizer(self.meshes_world, raster_settings=raster_settings)
+
 
 class TestPointRasterizer(unittest.TestCase):
     def test_simple_sphere(self):

tests/test_render_multigpu.py

@@ -14,6 +14,7 @@ from pytorch3d.renderer import (
     HardGouraudShader,
     Materials,
     MeshRasterizer,
+    MeshRasterizerOpenGL,
     MeshRenderer,
     PointLights,
     PointsRasterizationSettings,
@@ -21,18 +22,19 @@ from pytorch3d.renderer import (
     PointsRenderer,
     RasterizationSettings,
     SoftPhongShader,
+    SplatterPhongShader,
     TexturesVertex,
 )
 from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
 from pytorch3d.structures import Meshes, Pointclouds
 from pytorch3d.utils.ico_sphere import ico_sphere
 
-from .common_testing import get_random_cuda_device, TestCaseMixin
+from .common_testing import TestCaseMixin
 
 
 # Set the number of GPUS you want to test with
-NUM_GPUS = 3
-GPU_LIST = list({get_random_cuda_device() for _ in range(NUM_GPUS)})
+NUM_GPUS = 2
+GPU_LIST = [f"cuda:{idx}" for idx in range(NUM_GPUS)]
 print("GPUs: %s" % ", ".join(GPU_LIST))
 
 
@@ -56,12 +58,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
         self.assertEqual(renderer.shader.materials.device, device)
         self.assertEqual(renderer.shader.materials.ambient_color.device, device)
 
-    def test_mesh_renderer_to(self):
+    def _mesh_renderer_to(self, rasterizer_class, shader_class):
         """
         Test moving all the tensors in the mesh renderer to a new device.
         """
 
-        device1 = torch.device("cpu")
+        device1 = torch.device("cuda:0")
 
         R, T = look_at_view_transform(1500, 0.0, 0.0)
 
@@ -71,12 +73,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
         lights.location = torch.tensor([0.0, 0.0, +1000.0], device=device1)[None]
 
         raster_settings = RasterizationSettings(
-            image_size=256, blur_radius=0.0, faces_per_pixel=1
+            image_size=128, blur_radius=0.0, faces_per_pixel=1
         )
         cameras = FoVPerspectiveCameras(
             device=device1, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=100
         )
-        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
+        rasterizer = rasterizer_class(cameras=cameras, raster_settings=raster_settings)
 
         blend_params = BlendParams(
             1e-4,
@@ -84,7 +86,7 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
             background_color=torch.zeros(3, dtype=torch.float32, device=device1),
         )
 
-        shader = SoftPhongShader(
+        shader = shader_class(
             lights=lights,
             cameras=cameras,
             materials=materials,
@@ -107,26 +109,32 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
         # Move renderer and mesh to another device and re render
         # This also tests that background_color is correctly moved to
         # the new device
-        device2 = torch.device("cuda:0")
+        device2 = torch.device("cuda:1")
         renderer = renderer.to(device2)
         mesh = mesh.to(device2)
         self._check_mesh_renderer_props_on_device(renderer, device2)
         output_images = renderer(mesh)
         self.assertEqual(output_images.device, device2)
 
-    def test_render_meshes(self):
+    def test_mesh_renderer_to(self):
+        self._mesh_renderer_to(MeshRasterizer, SoftPhongShader)
+
+    def test_mesh_renderer_opengl_to(self):
+        self._mesh_renderer_to(MeshRasterizerOpenGL, SplatterPhongShader)
+
+    def _render_meshes(self, rasterizer_class, shader_class):
         test = self
 
         class Model(nn.Module):
-            def __init__(self):
+            def __init__(self, device):
                 super(Model, self).__init__()
-                mesh = ico_sphere(3)
+                mesh = ico_sphere(3).to(device)
                 self.register_buffer("faces", mesh.faces_padded())
-                self.renderer = self.init_render()
+                self.renderer = self.init_render(device)
 
-            def init_render(self):
+            def init_render(self, device):
 
-                cameras = FoVPerspectiveCameras()
+                cameras = FoVPerspectiveCameras().to(device)
                 raster_settings = RasterizationSettings(
                     image_size=128, blur_radius=0.0, faces_per_pixel=1
                 )
@@ -135,12 +143,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
                     diffuse_color=((0, 0.0, 0),),
                     specular_color=((0.0, 0, 0),),
                     location=((0.0, 0.0, 1e5),),
-                )
+                ).to(device)
                 renderer = MeshRenderer(
-                    rasterizer=MeshRasterizer(
+                    rasterizer=rasterizer_class(
                         cameras=cameras, raster_settings=raster_settings
                     ),
-                    shader=HardGouraudShader(cameras=cameras, lights=lights),
+                    shader=shader_class(cameras=cameras, lights=lights),
                 )
                 return renderer
 
@@ -155,20 +163,25 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
                 img_render = self.renderer(mesh)
                 return img_render[:, :, :, :3]
 
-        # DataParallel requires every input tensor be provided
-        # on the first device in its device_ids list.
-        verts = ico_sphere(3).verts_padded()
+        # Make sure we use all GPUs in GPU_LIST by making the batch size 4 x GPU count.
+        verts = ico_sphere(3).verts_padded().expand(len(GPU_LIST) * 4, 642, 3)
         texs = verts.new_ones(verts.shape)
-        model = Model()
-        model.to(GPU_LIST[0])
+        model = Model(device=GPU_LIST[0])
         model = nn.DataParallel(model, device_ids=GPU_LIST)
 
         # Test a few iterations
         for _ in range(100):
             model(verts, texs)
 
+    def test_render_meshes(self):
+        self._render_meshes(MeshRasterizer, HardGouraudShader)
 
-class TestRenderPointssMultiGPU(TestCaseMixin, unittest.TestCase):
+    # @unittest.skip("Multi-GPU OpenGL training is currently not supported.")
+    def test_render_meshes_opengl(self):
+        self._render_meshes(MeshRasterizerOpenGL, SplatterPhongShader)
+
+
+class TestRenderPointsMultiGPU(TestCaseMixin, unittest.TestCase):
     def _check_points_renderer_props_on_device(self, renderer, device):
         """
         Helper function to check that all the properties have