test_render_multigpu.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import unittest

import torch
import torch.nn as nn
from pytorch3d.renderer import (
    AlphaCompositor,
    BlendParams,
    HardGouraudShader,
    Materials,
    MeshRasterizer,
    MeshRenderer,
    PointLights,
    PointsRasterizationSettings,
    PointsRasterizer,
    PointsRenderer,
    RasterizationSettings,
    SoftPhongShader,
    SplatterPhongShader,
    TexturesVertex,
)
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
from pytorch3d.renderer.opengl import MeshRasterizerOpenGL
from pytorch3d.structures import Meshes, Pointclouds
from pytorch3d.utils.ico_sphere import ico_sphere

from .common_testing import TestCaseMixin


# Set the number of GPUS you want to test with
NUM_GPUS = 2
GPU_LIST = [f"cuda:{idx}" for idx in range(NUM_GPUS)]
print("GPUs: %s" % ", ".join(GPU_LIST))


class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
    def _check_mesh_renderer_props_on_device(self, renderer, device):
        """
        Helper function to check that all the properties of the mesh
        renderer have been moved to the correct device.
        """
        # Cameras
        self.assertEqual(renderer.rasterizer.cameras.device, device)
        self.assertEqual(renderer.rasterizer.cameras.R.device, device)
        self.assertEqual(renderer.rasterizer.cameras.T.device, device)
        self.assertEqual(renderer.shader.cameras.device, device)
        self.assertEqual(renderer.shader.cameras.R.device, device)
        self.assertEqual(renderer.shader.cameras.T.device, device)

        # Lights and Materials
        self.assertEqual(renderer.shader.lights.device, device)
        self.assertEqual(renderer.shader.lights.ambient_color.device, device)
        self.assertEqual(renderer.shader.materials.device, device)
        self.assertEqual(renderer.shader.materials.ambient_color.device, device)

    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("cuda:0")

        R, T = look_at_view_transform(1500, 0.0, 0.0)

        # Init shader settings
        materials = Materials(device=device1)
        lights = PointLights(device=device1)
        lights.location = torch.tensor([0.0, 0.0, +1000.0], device=device1)[None]

        raster_settings = RasterizationSettings(
            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 = rasterizer_class(cameras=cameras, raster_settings=raster_settings)

        blend_params = BlendParams(
            1e-4,
            1e-4,
            background_color=torch.zeros(3, dtype=torch.float32, device=device1),
        )

        shader = shader_class(
            lights=lights,
            cameras=cameras,
            materials=materials,
            blend_params=blend_params,
        )
        renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)

        mesh = ico_sphere(2, device1)
        verts_padded = mesh.verts_padded()
        textures = TexturesVertex(
            verts_features=torch.ones_like(verts_padded, device=device1)
        )
        mesh.textures = textures
        self._check_mesh_renderer_props_on_device(renderer, device1)

        # Test rendering on cpu
        output_images = renderer(mesh)
        self.assertEqual(output_images.device, device1)

        # 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: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_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, device):
                super(Model, self).__init__()
                mesh = ico_sphere(3).to(device)
                self.register_buffer("faces", mesh.faces_padded())
                self.renderer = self.init_render(device)

            def init_render(self, device):

                cameras = FoVPerspectiveCameras().to(device)
                raster_settings = RasterizationSettings(
                    image_size=128, blur_radius=0.0, faces_per_pixel=1
                )
                lights = PointLights(
                    ambient_color=((1.0, 1.0, 1.0),),
                    diffuse_color=((0, 0.0, 0),),
                    specular_color=((0.0, 0, 0),),
                    location=((0.0, 0.0, 1e5),),
                ).to(device)
                renderer = MeshRenderer(
                    rasterizer=rasterizer_class(
                        cameras=cameras, raster_settings=raster_settings
                    ),
                    shader=shader_class(cameras=cameras, lights=lights),
                )
                return renderer

            def forward(self, verts, texs):
                batch_size = verts.size(0)
                self.renderer = self.renderer.to(verts.device)
                tex = TexturesVertex(verts_features=texs)
                faces = self.faces.expand(batch_size, -1, -1).to(verts.device)
                mesh = Meshes(verts, faces, tex).to(verts.device)

                test._check_mesh_renderer_props_on_device(self.renderer, verts.device)
                img_render = self.renderer(mesh)
                return img_render[:, :, :, :3]

        # 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(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)

    # @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
        been moved to the correct device.
        """
        # Cameras
        self.assertEqual(renderer.rasterizer.cameras.device, device)
        self.assertEqual(renderer.rasterizer.cameras.R.device, device)
        self.assertEqual(renderer.rasterizer.cameras.T.device, device)

    def test_points_renderer_to(self):
        """
        Test moving all the tensors in the points renderer to a new device.
        """

        device1 = torch.device("cpu")

        R, T = look_at_view_transform(1500, 0.0, 0.0)

        raster_settings = PointsRasterizationSettings(
            image_size=256, radius=0.001, points_per_pixel=1
        )
        cameras = FoVPerspectiveCameras(
            device=device1, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=100
        )
        rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)

        renderer = PointsRenderer(rasterizer=rasterizer, compositor=AlphaCompositor())

        mesh = ico_sphere(2, device1)
        verts_padded = mesh.verts_padded()
        pointclouds = Pointclouds(
            points=verts_padded, features=torch.randn_like(verts_padded)
        )
        self._check_points_renderer_props_on_device(renderer, device1)

        # Test rendering on cpu
        output_images = renderer(pointclouds)
        self.assertEqual(output_images.device, device1)

        # Move renderer and pointclouds to another device and re render
        device2 = torch.device("cuda:0")
        renderer = renderer.to(device2)
        pointclouds = pointclouds.to(device2)
        self._check_points_renderer_props_on_device(renderer, device2)
        output_images = renderer(pointclouds)
        self.assertEqual(output_images.device, device2)