Spaces:
Running
Running
File size: 3,902 Bytes
7088d16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
# 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.
"""Test render speed."""
import logging
import sys
from os import path
import torch
from fvcore.common.benchmark import benchmark
from pytorch3d.renderer.points.pulsar import Renderer
from torch.autograd import Variable
# Making sure you can run this, even if pulsar hasn't been installed yet.
sys.path.insert(0, path.join(path.dirname(__file__), ".."))
LOGGER = logging.getLogger(__name__)
"""Measure the execution speed of the rendering.
This measures a very pessimistic upper bound on speed, because synchronization
points have to be introduced in Python. On a pure PyTorch execution pipeline,
results should be significantly faster. You can get pure CUDA timings through
C++ by activating `PULSAR_TIMINGS_BATCHED_ENABLED` in the file
`pytorch3d/csrc/pulsar/logging.h` or defining it for your compiler.
"""
def _bm_pulsar():
n_points = 1_000_000
width = 1_000
height = 1_000
renderer = Renderer(width, height, n_points)
# Generate sample data.
torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0
vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0
vert_col = torch.rand(n_points, 3, dtype=torch.float32)
vert_rad = torch.rand(n_points, dtype=torch.float32)
cam_params = torch.tensor(
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 2.0], dtype=torch.float32
)
device = torch.device("cuda")
vert_pos = vert_pos.to(device)
vert_col = vert_col.to(device)
vert_rad = vert_rad.to(device)
cam_params = cam_params.to(device)
renderer = renderer.to(device)
vert_pos_var = Variable(vert_pos, requires_grad=False)
vert_col_var = Variable(vert_col, requires_grad=False)
vert_rad_var = Variable(vert_rad, requires_grad=False)
cam_params_var = Variable(cam_params, requires_grad=False)
def bm_closure():
renderer.forward(
vert_pos_var,
vert_col_var,
vert_rad_var,
cam_params_var,
1.0e-1,
45.0,
percent_allowed_difference=0.01,
)
torch.cuda.synchronize()
return bm_closure
def _bm_pulsar_backward():
n_points = 1_000_000
width = 1_000
height = 1_000
renderer = Renderer(width, height, n_points)
# Generate sample data.
torch.manual_seed(1)
vert_pos = torch.rand(n_points, 3, dtype=torch.float32) * 10.0
vert_pos[:, 2] += 25.0
vert_pos[:, :2] -= 5.0
vert_col = torch.rand(n_points, 3, dtype=torch.float32)
vert_rad = torch.rand(n_points, dtype=torch.float32)
cam_params = torch.tensor(
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 2.0], dtype=torch.float32
)
device = torch.device("cuda")
vert_pos = vert_pos.to(device)
vert_col = vert_col.to(device)
vert_rad = vert_rad.to(device)
cam_params = cam_params.to(device)
renderer = renderer.to(device)
vert_pos_var = Variable(vert_pos, requires_grad=True)
vert_col_var = Variable(vert_col, requires_grad=True)
vert_rad_var = Variable(vert_rad, requires_grad=True)
cam_params_var = Variable(cam_params, requires_grad=True)
res = renderer.forward(
vert_pos_var,
vert_col_var,
vert_rad_var,
cam_params_var,
1.0e-1,
45.0,
percent_allowed_difference=0.01,
)
loss = res.sum()
def bm_closure():
loss.backward(retain_graph=True)
torch.cuda.synchronize()
return bm_closure
def bm_pulsar() -> None:
if not torch.cuda.is_available():
return
benchmark(_bm_pulsar, "PULSAR_FORWARD", [{}], warmup_iters=3)
benchmark(_bm_pulsar_backward, "PULSAR_BACKWARD", [{}], warmup_iters=3)
if __name__ == "__main__":
bm_pulsar()
|