def test_origin_map(self): manager = ME.CoordinateManager( D=1, coordinate_map_type=ME.CoordinateMapType.CPU) coords = torch.IntTensor([[0, -3], [0, -2], [0, -1], [0, 0], [1, 1], [1, 2], [1, 3]]) # key with batch_size 2 key, (unique_map, inverse_map) = manager.insert_and_map(coords, [1]) batch_indices, origin_map = manager.origin_map(key) print(origin_map) # self.assertTrue(set(origin_map[0].numpy()) == set([0, 1, 2, 3])) key = manager.origin() batch_coordinates = manager.get_coordinates(key) print(batch_coordinates) self.assertTrue(len(batch_coordinates) == 2) if not ME.is_cuda_available(): return manager = ME.CoordinateManager( D=1, coordinate_map_type=ME.CoordinateMapType.CUDA, allocator_type=ME.GPUMemoryAllocatorType.PYTORCH, ) key, (unique_map, inverse_map) = manager.insert_and_map(coords.to(0), [1]) origin_map = manager.origin_map(key) print(origin_map) key = manager.origin() self.assertTrue(manager.number_of_unique_batch_indices() == 2) batch_coordinates = manager.get_coordinates(key) print(batch_coordinates) self.assertTrue(len(batch_coordinates) == 2)
def test_gpu_allocator(self): if not ME.is_cuda_available(): return # Set the global GPU memory manager backend. By default PYTORCH. ME.set_gpu_allocator(ME.GPUMemoryAllocatorType.PYTORCH) ME.set_gpu_allocator(ME.GPUMemoryAllocatorType.CUDA) # Create a coords man with the specified GPU memory manager backend. # No effect with CPU_ONLY build manager = ME.CoordinateManager( D=1, coordinate_map_type=ME.CoordinateMapType.CPU, allocator_type=ME.GPUMemoryAllocatorType.CUDA, )
def test(self): self.assertTrue(ME.is_cuda_available() == torch.cuda.is_available()) if ME.is_cuda_available(): print(ME.cuda_version()) print(ME.get_gpu_memory_info())