| 16 | |
| 17 | class TestLimit(unittest.TestCase): |
| 18 | def _test( |
| 19 | self, |
| 20 | b: int = 1, |
| 21 | n: int = 10000, # number of points |
| 22 | m: int = 20000, # number of rays |
| 23 | k: int = 40, # number of neighboring points |
| 24 | ray_radius: float = 0.1, |
| 25 | grid_size: int = 100, |
| 26 | grid_width: float = 1., |
| 27 | ): |
| 28 | print('very beginning') |
| 29 | torch.cuda.empty_cache() |
| 30 | print(f'cuda memory reserved: {torch.cuda.memory_reserved(0)/1.e6} MB') |
| 31 | print(f'cuda memory allocated: {torch.cuda.memory_allocated(0)/1.e6} MB') |
| 32 | |
| 33 | points = (torch.rand(b, n, 3) - 0.5) * 2 * grid_width |
| 34 | ray_origins = torch.randn(b, m, 3) |
| 35 | ray_directions = torch.nn.functional.normalize(torch.randn(b, m, 3), dim=-1) |
| 36 | ray_radius = torch.ones(b) * ray_radius |
| 37 | grid_size = torch.ones(b, 3, dtype=torch.long) * grid_size |
| 38 | grid_width = torch.ones(b, 3) * grid_width |
| 39 | |
| 40 | |
| 41 | if not torch.cuda.is_available(): |
| 42 | return |
| 43 | |
| 44 | # test gpu |
| 45 | device = torch.device('cuda') |
| 46 | points = points.to(device=device) |
| 47 | ray_origins = ray_origins.to(device=device) |
| 48 | ray_directions = ray_directions.to(device=device) |
| 49 | ray_radius = ray_radius.to(device=device) if isinstance(ray_radius, torch.Tensor) else ray_radius |
| 50 | grid_size = grid_size.to(device=device) if isinstance(grid_size, torch.Tensor) else grid_size |
| 51 | grid_width = grid_width.to(device=device) if isinstance(grid_width, torch.Tensor) else grid_width |
| 52 | |
| 53 | print('send to gpu') |
| 54 | torch.cuda.empty_cache() |
| 55 | print(f'cuda memory reserved: {torch.cuda.memory_reserved(0) / 1.e6} MB') |
| 56 | print(f'cuda memory allocated: {torch.cuda.memory_allocated(0) / 1.e6} MB') |
| 57 | |
| 58 | stime = timer() |
| 59 | all_ray2pidxs = pr_utils.find_k_neighbor_points_of_rays( |
| 60 | points=points, |
| 61 | k=k, |
| 62 | ray_origins=ray_origins, |
| 63 | ray_directions=ray_directions, |
| 64 | ray_radius=ray_radius, |
| 65 | grid_size=grid_size, |
| 66 | grid_width=grid_width, |
| 67 | version='v2', |
| 68 | ) |
| 69 | time_pr_k_cuda1 = timer() - stime |
| 70 | print(f'pr_k_cuda_1: {time_pr_k_cuda1:.3f} secs') |
| 71 | |
| 72 | torch.cuda.empty_cache() |
| 73 | print('after find_k_neighbor_points_of_rays') |
| 74 | print(f'cuda memory reserved: {torch.cuda.memory_reserved(0) / 1.e6} MB') |
| 75 | print(f'cuda memory allocated: {torch.cuda.memory_allocated(0) / 1.e6} MB') |