pytorch3d/tests/bm_points_alignment.py
Jeremy Reizenstein b87058c62a fix recent lint
Summary: lint clean again

Reviewed By: patricklabatut

Differential Revision: D20868775

fbshipit-source-id: ade4301c1012c5c6943186432465215701d635a9
2020-04-06 06:41:00 -07:00

42 lines
1.2 KiB
Python

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
from copy import deepcopy
from itertools import product
from fvcore.common.benchmark import benchmark
from test_points_alignment import TestCorrespondingPointsAlignment
def bm_corresponding_points_alignment() -> None:
case_grid = {
"allow_reflection": [True, False],
"batch_size": [1, 10, 100],
"dim": [3, 20],
"estimate_scale": [True, False],
"n_points": [100, 10000],
"random_weights": [False, True],
"use_pointclouds": [False],
}
test_args = sorted(case_grid.keys())
test_cases = product(*[case_grid[k] for k in test_args])
kwargs_list = [dict(zip(test_args, case)) for case in test_cases]
# add the use_pointclouds=True test cases whenever we have dim==3
kwargs_to_add = []
for entry in kwargs_list:
if entry["dim"] == 3:
entry_add = deepcopy(entry)
entry_add["use_pointclouds"] = True
kwargs_to_add.append(entry_add)
kwargs_list.extend(kwargs_to_add)
benchmark(
TestCorrespondingPointsAlignment.corresponding_points_alignment,
"CorrespodingPointsAlignment",
kwargs_list,
warmup_iters=1,
)