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())