Skip to content

Commit ee2b2fe

Browse files
bottlerfacebook-github-bot
bottler
authored and
facebook-github-bot
committed
Use C++/CUDA in points2vols
Summary: Move the core of add_points_to_volumes to the new C++/CUDA implementation. Add new flag to let the user stop this happening. Avoids copies. About a 30% speedup on the larger cases, up to 50% on the smaller cases. New timings ``` Benchmark Avg Time(μs) Peak Time(μs) Iterations -------------------------------------------------------------------------------- ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_1000 4575 12591 110 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_10000 25468 29186 20 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_100000 202085 209897 3 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_1000 46059 48188 11 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_10000 83759 95669 7 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_100000 326056 339393 2 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_1000 2379 4738 211 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_10000 12100 63099 42 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_100000 63323 63737 8 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_1000 45216 45479 12 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_10000 57205 58524 9 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_100000 139499 139926 4 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_1000 40129 40431 13 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_10000 204949 239293 3 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_100000 1664541 1664541 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_1000 391573 395108 2 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_10000 674869 674869 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_100000 2713632 2713632 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_1000 12726 13506 40 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_10000 73103 73299 7 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_100000 598634 598634 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_1000 398742 399256 2 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_10000 543129 543129 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_100000 1242956 1242956 1 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_1000 1814 8884 276 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_10000 1996 8851 251 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_100000 4608 11529 109 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_1000 5183 12508 97 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_10000 7106 14077 71 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_100000 25914 31818 20 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_1000 1778 8823 282 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_10000 1825 8613 274 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_100000 3154 10161 159 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_1000 4888 9404 103 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_10000 5194 9963 97 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_100000 8109 14933 62 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_1000 3320 10306 151 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_10000 7003 8595 72 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_100000 49140 52957 11 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_1000 35890 36918 14 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_10000 58890 59337 9 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_100000 286878 287600 2 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_1000 2484 8805 202 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_10000 3967 9090 127 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_100000 19423 19799 26 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_1000 33228 33329 16 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_10000 37292 37370 14 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_100000 73550 74017 7 -------------------------------------------------------------------------------- ``` Previous timings ``` Benchmark Avg Time(μs) Peak Time(μs) Iterations -------------------------------------------------------------------------------- ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_1000 10100 46422 50 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_10000 28442 32100 18 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[25, 25, 25]_100000 241127 254269 3 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_1000 54149 79480 10 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_10000 125459 212734 4 ADD_POINTS_TO_VOLUMES_cpu_10_trilinear_[101, 111, 121]_100000 512739 512739 1 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_1000 2866 13365 175 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_10000 7026 12604 72 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[25, 25, 25]_100000 48822 55607 11 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_1000 38098 38576 14 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_10000 48006 54120 11 ADD_POINTS_TO_VOLUMES_cpu_10_nearest_[101, 111, 121]_100000 131563 138536 4 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_1000 64615 91735 8 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_10000 228815 246095 3 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[25, 25, 25]_100000 3086615 3086615 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_1000 464298 465292 2 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_10000 1053440 1053440 1 ADD_POINTS_TO_VOLUMES_cpu_100_trilinear_[101, 111, 121]_100000 6736236 6736236 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_1000 11940 12440 42 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_10000 56641 58051 9 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[25, 25, 25]_100000 711492 711492 1 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_1000 326437 329846 2 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_10000 418514 427911 2 ADD_POINTS_TO_VOLUMES_cpu_100_nearest_[101, 111, 121]_100000 1524285 1524285 1 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_1000 5949 13602 85 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_10000 5817 13001 86 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[25, 25, 25]_100000 23833 25971 21 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_1000 9029 16178 56 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_10000 11595 18601 44 ADD_POINTS_TO_VOLUMES_cuda:0_10_trilinear_[101, 111, 121]_100000 46986 47344 11 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_1000 2554 9747 196 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_10000 2676 9537 187 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[25, 25, 25]_100000 6567 14179 77 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_1000 5840 12811 86 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_10000 6102 13128 82 ADD_POINTS_TO_VOLUMES_cuda:0_10_nearest_[101, 111, 121]_100000 11945 11995 42 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_1000 7642 13671 66 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_10000 25190 25260 20 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[25, 25, 25]_100000 212018 212134 3 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_1000 40421 45692 13 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_10000 92078 92132 6 ADD_POINTS_TO_VOLUMES_cuda:0_100_trilinear_[101, 111, 121]_100000 457211 457229 2 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_1000 3574 10377 140 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_10000 7222 13023 70 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[25, 25, 25]_100000 48127 48165 11 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_1000 34732 35295 15 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_10000 43050 51064 12 ADD_POINTS_TO_VOLUMES_cuda:0_100_nearest_[101, 111, 121]_100000 106028 106058 5 -------------------------------------------------------------------------------- ``` Reviewed By: nikhilaravi Differential Revision: D29548609 fbshipit-source-id: 7026e832ea299145c3f6b55687f3c1601294f5c0
1 parent 9ad98c8 commit ee2b2fe

File tree

2 files changed

+76
-7
lines changed

2 files changed

+76
-7
lines changed

‎pytorch3d/ops/points_to_volumes.py

+70-5
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ def add_pointclouds_to_volumes(
192192
initial_volumes: "Volumes",
193193
mode: str = "trilinear",
194194
min_weight: float = 1e-4,
195+
_python: bool = False,
195196
) -> "Volumes":
196197
"""
197198
Add a batch of point clouds represented with a `Pointclouds` structure
@@ -249,6 +250,8 @@ def add_pointclouds_to_volumes(
249250
min_weight: A scalar controlling the lowest possible total per-voxel
250251
weight used to normalize the features accumulated in a voxel.
251252
Only active for `mode==trilinear`.
253+
_python: Set to True to use a pure Python implementation, e.g. for test
254+
purposes, which requires more memory and may be slower.
252255
253256
Returns:
254257
updated_volumes: Output `Volumes` structure containing the conversion result.
@@ -283,6 +286,7 @@ def add_pointclouds_to_volumes(
283286
grid_sizes=initial_volumes.get_grid_sizes(),
284287
mask=mask,
285288
mode=mode,
289+
_python=_python,
286290
)
287291

288292
return initial_volumes.update_padded(
@@ -299,6 +303,7 @@ def add_points_features_to_volume_densities_features(
299303
min_weight: float = 1e-4,
300304
mask: Optional[torch.Tensor] = None,
301305
grid_sizes: Optional[torch.LongTensor] = None,
306+
_python: bool = False,
302307
) -> Tuple[torch.Tensor, torch.Tensor]:
303308
"""
304309
Convert a batch of point clouds represented with tensors of per-point
@@ -340,6 +345,7 @@ def add_points_features_to_volume_densities_features(
340345
grid_sizes: `LongTensor` of shape (minibatch, 3) representing the
341346
spatial resolutions of each of the the non-flattened `volumes` tensors,
342347
or None to indicate the whole volume is used for every batch element.
348+
_python: Set to True to use a pure Python implementation.
343349
Returns:
344350
volume_features: Output volume of shape `(minibatch, feature_dim, D, H, W)`
345351
volume_densities: Occupancy volume of shape `(minibatch, 1, D, H, W)`
@@ -362,6 +368,66 @@ def add_points_features_to_volume_densities_features(
362368
.expand(volume_densities.shape[0], 3)
363369
)
364370

371+
if _python:
372+
return _add_points_features_to_volume_densities_features_python(
373+
points_3d=points_3d,
374+
points_features=points_features,
375+
volume_densities=volume_densities,
376+
volume_features=volume_features,
377+
mode=mode,
378+
min_weight=min_weight,
379+
mask=mask,
380+
grid_sizes=grid_sizes,
381+
)
382+
383+
if mode == "trilinear":
384+
splat = True
385+
elif mode == "nearest":
386+
splat = False
387+
else:
388+
raise ValueError('No such interpolation mode "%s"' % mode)
389+
volume_densities, volume_features = _points_to_volumes(
390+
points_3d,
391+
points_features,
392+
volume_densities,
393+
volume_features,
394+
grid_sizes,
395+
1.0, # point_weight
396+
mask,
397+
True, # align_corners
398+
splat,
399+
)
400+
if splat:
401+
# divide each feature by the total weight of the votes
402+
volume_features = volume_features / volume_densities.clamp(min_weight)
403+
else:
404+
# divide each feature by the total weight of the votes
405+
volume_features = volume_features / volume_densities.clamp(1.0)
406+
407+
return volume_features, volume_densities
408+
409+
410+
def _add_points_features_to_volume_densities_features_python(
411+
*,
412+
points_3d: torch.Tensor,
413+
points_features: torch.Tensor,
414+
volume_densities: torch.Tensor,
415+
volume_features: Optional[torch.Tensor],
416+
mode: str,
417+
min_weight: float,
418+
mask: Optional[torch.Tensor],
419+
grid_sizes: torch.LongTensor,
420+
) -> Tuple[torch.Tensor, torch.Tensor]:
421+
"""
422+
Python implementation for add_points_features_to_volume_densities_features.
423+
424+
Returns:
425+
volume_features: Output volume of shape `(minibatch, feature_dim, D, H, W)`
426+
volume_densities: Occupancy volume of shape `(minibatch, 1, D, H, W)`
427+
containing the total amount of votes cast to each of the voxels.
428+
"""
429+
ba, n_points, feature_dim = points_features.shape
430+
365431
# flatten densities and features
366432
v_shape = volume_densities.shape[2:]
367433
volume_densities_flatten = volume_densities.view(ba, -1, 1)
@@ -376,7 +442,7 @@ def add_points_features_to_volume_densities_features(
376442
volume_features_flatten = volume_features.view(ba, feature_dim, n_voxels)
377443

378444
if mode == "trilinear": # do the splatting (trilinear interp)
379-
volume_features, volume_densities = splat_points_to_volumes(
445+
volume_features, volume_densities = _splat_points_to_volumes(
380446
points_3d,
381447
points_features,
382448
volume_densities_flatten,
@@ -386,7 +452,7 @@ def add_points_features_to_volume_densities_features(
386452
min_weight=min_weight,
387453
)
388454
elif mode == "nearest": # nearest neighbor interp
389-
volume_features, volume_densities = round_points_to_volumes(
455+
volume_features, volume_densities = _round_points_to_volumes(
390456
points_3d,
391457
points_features,
392458
volume_densities_flatten,
@@ -400,7 +466,6 @@ def add_points_features_to_volume_densities_features(
400466
# reshape into the volume shape
401467
volume_features = volume_features.view(ba, feature_dim, *v_shape)
402468
volume_densities = volume_densities.view(ba, 1, *v_shape)
403-
404469
return volume_features, volume_densities
405470

406471

@@ -441,7 +506,7 @@ def _check_points_to_volumes_inputs(
441506
)
442507

443508

444-
def splat_points_to_volumes(
509+
def _splat_points_to_volumes(
445510
points_3d: torch.Tensor,
446511
points_features: torch.Tensor,
447512
volume_densities: torch.Tensor,
@@ -574,7 +639,7 @@ def splat_points_to_volumes(
574639
return volume_features, volume_densities
575640

576641

577-
def round_points_to_volumes(
642+
def _round_points_to_volumes(
578643
points_3d: torch.Tensor,
579644
points_features: torch.Tensor,
580645
volume_densities: torch.Tensor,

‎tests/test_points_to_volumes.py

+6-2
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66

77
import unittest
88
from functools import partial
9+
from itertools import product
910
from typing import Tuple
1011

1112
import numpy as np
@@ -254,7 +255,7 @@ def test_from_point_cloud(self, interp_mode="trilinear"):
254255

255256
for volume_size in ([25, 25, 25], [30, 25, 15]):
256257

257-
for interp_mode in ("trilinear", "nearest"):
258+
for python, interp_mode in product([True, False], ["trilinear", "nearest"]):
258259

259260
(pointclouds, initial_volumes) = init_volume_boundary_pointcloud(
260261
volume_size=volume_size,
@@ -266,7 +267,10 @@ def test_from_point_cloud(self, interp_mode="trilinear"):
266267
)
267268

268269
volumes = add_pointclouds_to_volumes(
269-
pointclouds, initial_volumes, mode=interp_mode
270+
pointclouds,
271+
initial_volumes,
272+
mode=interp_mode,
273+
_python=python,
270274
)
271275

272276
V_color, V_density = volumes.features(), volumes.densities()

0 commit comments

Comments
 (0)