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_stride(self):
coordinates = torch.IntTensor([[0, 1], [0, 1], [0, 2], [0, 2], [1, 0],
[1, 0], [1, 1]])
manager = ME.CoordinateManager(
D=1, coordinate_map_type=ME.CoordinateMapType.CPU)
key, (unique_map,
inverse_map) = manager.insert_and_map(coordinates, [1])
# Create a strided map
stride_key = manager.stride(key, [4])
print(manager.get_coordinates(key))
print(manager.get_coordinates(stride_key))
print(
manager.kernel_map(
key,
stride_key,
[4],
[4],
[1],
ME.RegionType.HYPER_CUBE,
torch.IntTensor(),
False,
True,
))
def test_coordinate_manager(self):
coordinates = torch.IntTensor([[0, 1], [0, 1], [0, 2], [0, 2], [1, 0],
[1, 0], [1, 1]])
manager = ME.CoordinateManager(
D=1, coordinate_map_type=ME.CoordinateMapType.CPU)
key, (unique_map,
inverse_map) = manager.insert_and_map(coordinates, [1])
# mapping and inverse mapping should recover the original coordinates
self.assertTrue(
torch.all(coordinates[unique_map.long()][inverse_map.long()] ==
coordinates))
# copied coordinates should retrieve the original coordinates
retrieved_coordinates = manager.get_coordinates(key)
self.assertTrue(
torch.all(
coordinates == retrieved_coordinates[inverse_map.long()]))
# Create a strided map
stride_key = manager.stride(key, [4])
strided_coords = manager.get_coordinates(stride_key)
self.assertTrue(len(strided_coords) == 2)
def test_stride_cuda(self):
coordinates = torch.IntTensor([[0, 1], [0, 1], [0, 2], [0, 2], [1, 0],
[1, 0], [1, 1]]).cuda()
manager = ME.CoordinateManager(
D=1, coordinate_map_type=ME.CoordinateMapType.CUDA)
key, (unique_map,
inverse_map) = manager.insert_and_map(coordinates, [1])
# Create a strided map
stride_key = manager.stride(key, [4])
print(manager.get_coordinates(key))
print(manager.get_coordinates(stride_key))
# print(
# manager.kernel_map(
# key,
# stride_key,
# [4],
# [4],
# [1],
# ME.RegionType.HYPER_CUBE,
# torch.IntTensor(),
# False,
# True,
# )
# )
print(manager.stride_map(key, stride_key))
def quantize(coordinates):
D = coordinates.size(1) - 1
coordinate_manager = ME.CoordinateManager(
D=D, coordinate_map_type=ME.CoordinateMapType.CPU)
coordinate_map_key = ME.CoordinateMapKey(convert_to_int_list(1, D), "")
key, (unique_map, inverse_map) = coordinate_manager.insert_and_map(
coordinates, *coordinate_map_key.get_key())
return unique_map, inverse_map
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_negative_coords(self):
coords = torch.IntTensor([[0, -3], [0, -2], [0, -1], [0, 0], [0, 1],
[0, 2], [0, 3]])
# Initialize map
manager = ME.CoordinateManager(
D=1, coordinate_map_type=ME.CoordinateMapType.CPU)
key, (unique_map, inverse_map) = manager.insert_and_map(coords, [1])
# Create a strided map
stride_key = manager.stride(key, [2])
strided_coords = manager.get_coordinates(stride_key).numpy().tolist()
self.assertTrue(len(strided_coords) == 4)
self.assertTrue([0, -4] in strided_coords)
self.assertTrue([0, -2] in strided_coords)
self.assertTrue([0, 2] in strided_coords)
