-
Notifications
You must be signed in to change notification settings - Fork 1.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Summary: Small fix and updated tests for multigpu rendering case. This resolves the issue seen in: #401 Reviewed By: gkioxari Differential Revision: D24314681 fbshipit-source-id: 84c5a5359844c77518b48044001daa9a86f3c43a
- Loading branch information
1 parent
4d52f9f
commit 563d441
Showing
3 changed files
with
162 additions
and
66 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,159 @@ | ||
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. | ||
|
||
import unittest | ||
|
||
import torch | ||
import torch.nn as nn | ||
from common_testing import TestCaseMixin, get_random_cuda_device | ||
from pytorch3d.renderer import ( | ||
BlendParams, | ||
HardGouraudShader, | ||
Materials, | ||
MeshRasterizer, | ||
MeshRenderer, | ||
PointLights, | ||
RasterizationSettings, | ||
SoftPhongShader, | ||
TexturesVertex, | ||
) | ||
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform | ||
from pytorch3d.structures.meshes import Meshes | ||
from pytorch3d.utils.ico_sphere import ico_sphere | ||
|
||
|
||
# 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)}) | ||
print("GPUs: %s" % ", ".join(GPU_LIST)) | ||
|
||
|
||
class TestRenderMultiGPU(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 test_mesh_renderer_to(self): | ||
""" | ||
Test moving all the tensors in the mesh renderer to a new device. | ||
""" | ||
|
||
device1 = torch.device("cpu") | ||
|
||
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=256, 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) | ||
|
||
blend_params = BlendParams( | ||
1e-4, | ||
1e-4, | ||
background_color=torch.zeros(3, dtype=torch.float32, device=device1), | ||
) | ||
|
||
shader = SoftPhongShader( | ||
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:0") | ||
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): | ||
test = self | ||
|
||
class Model(nn.Module): | ||
def __init__(self): | ||
super(Model, self).__init__() | ||
mesh = ico_sphere(3) | ||
self.register_buffer("faces", mesh.faces_padded()) | ||
self.renderer = self.init_render() | ||
|
||
def init_render(self): | ||
|
||
cameras = FoVPerspectiveCameras() | ||
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),), | ||
) | ||
renderer = MeshRenderer( | ||
rasterizer=MeshRasterizer( | ||
cameras=cameras, raster_settings=raster_settings | ||
), | ||
shader=HardGouraudShader(cameras=cameras, lights=lights), | ||
) | ||
return renderer | ||
|
||
def forward(self, verts, texs): | ||
batch_size = verts.size(0) | ||
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] | ||
|
||
# DataParallel requires every input tensor be provided | ||
# on the first device in its device_ids list. | ||
verts = ico_sphere(3).verts_padded() | ||
texs = verts.new_ones(verts.shape) | ||
model = Model() | ||
model = nn.DataParallel(model, device_ids=GPU_LIST) | ||
model.to(f"cuda:{model.device_ids[0]}") | ||
|
||
# Test a few iterations | ||
for _ in range(100): | ||
model(verts, texs) |