(self, b=10, n=20, max_size=20, mode='ind')
| 21 | class TestGetGridIdx(unittest.TestCase): |
| 22 | |
| 23 | def _test(self, b=10, n=20, max_size=20, mode='ind'): |
| 24 | if not torch.cuda.is_available(): |
| 25 | return |
| 26 | |
| 27 | points = torch.randn(b, n, 3) |
| 28 | grid_size = (torch.rand(b, 3) * (max_size-1) + 1).floor().long() |
| 29 | grid_center = torch.randn(b, 3) |
| 30 | grid_width = torch.rand(b, 3) * 2 + 0.1 |
| 31 | |
| 32 | stime = timer() |
| 33 | grid_idxs_gt, valid_mask_gt = naive.get_grid_idx( |
| 34 | points=points, |
| 35 | grid_size=grid_size, |
| 36 | center=grid_center, |
| 37 | grid_width=grid_width, |
| 38 | mode=mode, |
| 39 | ) |
| 40 | time_python = timer() - stime |
| 41 | |
| 42 | stime = timer() |
| 43 | grid_idxs, valid_mask = pr_cpp.get_grid_idx( |
| 44 | points, |
| 45 | grid_size, |
| 46 | grid_center, |
| 47 | grid_width, |
| 48 | mode, |
| 49 | ) |
| 50 | time_cpp = timer() - stime |
| 51 | |
| 52 | points = points.cuda() |
| 53 | grid_size = grid_size.cuda() |
| 54 | grid_center = grid_center.cuda() |
| 55 | grid_width = grid_width.cuda() |
| 56 | |
| 57 | stime = timer() |
| 58 | grid_idxs, valid_mask, cell_counts = pr_cuda.get_grid_idx( |
| 59 | points, |
| 60 | grid_size, |
| 61 | grid_center, |
| 62 | grid_width, |
| 63 | mode, |
| 64 | ) |
| 65 | time_cuda = timer() - stime |
| 66 | grid_idxs = grid_idxs.cpu() |
| 67 | valid_mask = valid_mask.cpu() |
| 68 | cell_counts = cell_counts.cpu() |
| 69 | |
| 70 | print('get_grid_idx:') |
| 71 | print(f'python: {time_python * 1000.:.3f} ms') |
| 72 | print(f'cpp: {time_cpp * 1000.:.3f} ms ({time_python / time_cpp:.2f} speed up)') |
| 73 | print(f'cuda: {time_cuda * 1000.:.3f} ms ({time_python / time_cuda:.2f} speed up)') |
| 74 | |
| 75 | assert torch.allclose(grid_idxs_gt, grid_idxs) |
| 76 | assert torch.allclose(valid_mask_gt, valid_mask) |
| 77 | |
| 78 | |
| 79 | # check cell count |
| 80 | batch_size = points.size(0) |
no test coverage detected