pytorch3d/tests/bm_raysampling.py
Jeremy Reizenstein 4711665edb lint
Summary: Fix recent lint.

Reviewed By: nikhilaravi

Differential Revision: D25900168

fbshipit-source-id: 6b6e8d35b68c8415ef305dc4719f43eda9316c8f
2021-01-20 13:08:35 -08:00

40 lines
1.1 KiB
Python

# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import itertools
from fvcore.common.benchmark import benchmark
from pytorch3d.renderer import (
FoVOrthographicCameras,
FoVPerspectiveCameras,
GridRaysampler,
MonteCarloRaysampler,
NDCGridRaysampler,
OrthographicCameras,
PerspectiveCameras,
)
from test_raysampling import TestRaysampling
def bm_raysampling() -> None:
case_grid = {
"raysampler_type": [GridRaysampler, NDCGridRaysampler, MonteCarloRaysampler],
"camera_type": [
PerspectiveCameras,
OrthographicCameras,
FoVPerspectiveCameras,
FoVOrthographicCameras,
],
"batch_size": [1, 10],
"n_pts_per_ray": [10, 1000, 10000],
"image_width": [10, 300],
"image_height": [10, 300],
}
test_cases = itertools.product(*case_grid.values())
kwargs_list = [dict(zip(case_grid.keys(), case)) for case in test_cases]
benchmark(TestRaysampling.raysampler, "RAYSAMPLER", kwargs_list, warmup_iters=1)
if __name__ == "__main__":
bm_raysampling()