def __init__( self, kernel_size=-1, stride=1, dilation=1, is_transpose: bool = False, region_type: RegionType = RegionType.HYPER_CUBE, region_offsets: torch.Tensor = None, expand_coordinates: bool = False, axis_types=None, dimension=-1, ): r""" :attr:`region_type` (RegionType, optional): defines the kernel shape. Please refer to MinkowskiEngine.Comon for details. :attr:`region_offset` (torch.IntTensor, optional): when the :attr:`region_type` is :attr:`RegionType.CUSTOM`, the convolution kernel uses the provided `region_offset` to define offsets. It should be a matrix of size :math:`N \times D` where :math:`N` is the number of offsets and :math:`D` is the dimension of the space. :attr:`axis_types` (list of RegionType, optional): If given, it uses different methods to create a kernel for each axis. e.g., when it is `[RegionType.HYPER_CUBE, RegionType.HYPER_CUBE, RegionType.HYPER_CROSS]`, the kernel would be rectangular for the first two dimensions and cross shaped for the thrid dimension. """ assert dimension > 0 assert isinstance(region_type, RegionType) kernel_size = convert_to_int_list(kernel_size, dimension) kernel_stride = convert_to_int_list(stride, dimension) kernel_dilation = convert_to_int_list(dilation, dimension) self.cache = {} self.kernel_size = kernel_size self.kernel_stride = kernel_stride self.kernel_dilation = kernel_dilation self.region_type = region_type self.region_offsets = region_offsets if region_offsets else torch.IntTensor( ) self.axis_types = axis_types self.dimension = dimension self.kernel_volume = get_kernel_volume(region_type, kernel_size, region_offsets, axis_types, dimension) self.requires_strided_coordinates = reduce( lambda s1, s2: s1 == 1 and s2 == 1, kernel_stride) self.expand_coordinates = expand_coordinates
def __init__( self, kernel_size, kernel_stride, kernel_dilation, region_type, offset, dimension, ): kernel_size = convert_to_int_list(kernel_size, dimension) kernel_stride = convert_to_int_list(kernel_stride, dimension) kernel_dilation = convert_to_int_list(kernel_dilation, dimension) super(KernelRegion, self).__init__(kernel_size, kernel_stride, kernel_dilation, region_type, offset, dimension)
def _get_coordinate_map_key( input: SparseTensor, coordinates: torch.Tensor = None, tensor_stride: StrideType = 1, expand_coordinates: bool = False, ): r"""Returns the coordinates map key.""" if coordinates is not None and not expand_coordinates: assert isinstance(coordinates, (CoordinateMapKey, torch.Tensor, SparseTensor)) if isinstance(coordinates, torch.Tensor): assert coordinates.ndim == 2 coordinate_map_key = CoordinateMapKey( convert_to_int_list(tensor_stride, coordinates.size(1) - 1), "") ( coordinate_map_key, (unique_index, inverse_mapping), ) = input._manager.insert_and_map(coordinates, *coordinate_map_key.get_key()) elif isinstance(coordinates, SparseTensor): coordinate_map_key = coordinates.coordinate_map_key else: # CoordinateMapKey type due to the previous assertion coordinate_map_key = coordinates else: # coordinates is None coordinate_map_key = CoordinateMapKey( input.coordinate_map_key.get_coordinate_size()) return coordinate_map_key
def insert_and_map( self, coordinates: torch.Tensor, tensor_stride: Union[int, Sequence, np.ndarray] = 1, string_id: str = "", ) -> Tuple[CoordinateMapKey, Tuple[torch.IntTensor, torch.IntTensor]]: r"""create a new coordinate map and returns (key, (map, inverse_map)). :attr:`coordinates`: `torch.Tensor` (Int tensor. `CUDA` if coordinate_map_type == `CoordinateMapType.GPU`) that defines the coordinates. :attr:`tensor_stride` (`list`): a list of `D` elements that defines the tensor stride for the new order-`D + 1` sparse tensor. Example:: >>> manager = CoordinateManager(D=1) >>> coordinates = torch.IntTensor([[0, 0], [0, 0], [0, 1], [0, 2]]) >>> key, (unique_map, inverse_map) = manager.insert(coordinates, [1]) >>> print(key) # key is tensor_stride, string_id [1]:"" >>> torch.all(coordinates[unique_map] == manager.get_coordinates(key)) # True >>> torch.all(coordinates == coordinates[unique_map][inverse_map]) # True """ tensor_stride = convert_to_int_list(tensor_stride, self.D) return self._manager.insert_and_map(coordinates, tensor_stride, string_id)
def unique_coordinate_map( coordinates: torch.Tensor, tensor_stride: Union[int, Sequence, np.ndarray] = 1, ) -> Tuple[torch.IntTensor, torch.IntTensor]: r"""Returns the unique indices and the inverse indices of the coordinates. :attr:`coordinates`: `torch.Tensor` (Int tensor. `CUDA` if coordinate_map_type == `CoordinateMapType.GPU`) that defines the coordinates. Example:: >>> coordinates = torch.IntTensor([[0, 0], [0, 0], [0, 1], [0, 2]]) >>> unique_map, inverse_map = unique_coordinates_map(coordinates) >>> coordinates[unique_map] # unique coordinates >>> torch.all(coordinates == coordinates[unique_map][inverse_map]) # True """ assert coordinates.ndim == 2, "Coordinates must be a matrix" assert isinstance(coordinates, torch.Tensor) if not coordinates.is_cuda: manager = MEB.CoordinateMapManagerCPU() else: manager = MEB.CoordinateMapManagerGPU_c10() tensor_stride = convert_to_int_list(tensor_stride, coordinates.shape[-1] - 1) key, (unique_map, inverse_map) = manager.insert_and_map(coordinates, tensor_stride, "") return unique_map, inverse_map
def kernel_map( self, in_key: CoordinateMapKey, out_key: CoordinateMapKey, stride=1, kernel_size=3, dilation=1, region_type=RegionType.HYPER_CUBE, region_offset=None, is_transpose=False, is_pool=False, ) -> dict: r"""Get kernel in-out maps for the specified coords keys or tensor strides. returns dict{kernel_index: in_out_tensor} where in_out_tensor[0] is the input row indices that correspond to in_out_tensor[1], which is the row indices for output. """ # region type 1 iteration with kernel_size 1 is invalid if isinstance(kernel_size, torch.Tensor): assert (kernel_size > 0).all(), f"Invalid kernel size: {kernel_size}" if (kernel_size == 1).all() == 1: region_type = RegionType.HYPER_CUBE elif isinstance(kernel_size, int): assert kernel_size > 0, f"Invalid kernel size: {kernel_size}" if kernel_size == 1: region_type = RegionType.HYPER_CUBE in_key = self._get_coordinate_map_key(in_key) out_key = self._get_coordinate_map_key(out_key) if region_offset is None: region_offset = torch.IntTensor() kernel_map = self._manager.kernel_map( in_key, out_key, convert_to_int_list(stride, self.D), # convert_to_int_list(kernel_size, self.D), # convert_to_int_list(dilation, self.D), # region_type, region_offset, is_transpose, is_pool, ) return kernel_map
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 sparse(self, tensor_stride: Union[int, Sequence, np.array] = 1, quantization_mode=None): r"""Converts the current sparse tensor field to a sparse tensor.""" if quantization_mode is None: quantization_mode = self.quantization_mode tensor_stride = convert_to_int_list(tensor_stride, self.D) sparse_tensor_key, ( unique_index, inverse_mapping, ) = self._manager.field_to_sparse_insert_and_map( self.coordinate_field_map_key, tensor_stride, ) self._inverse_mapping[sparse_tensor_key] = inverse_mapping if self.quantization_mode in [ SparseTensorQuantizationMode.UNWEIGHTED_SUM, SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE, ]: spmm = MinkowskiSPMMFunction() N = len(self._F) cols = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) vals = torch.ones(N, dtype=self._F.dtype, device=self._F.device) size = torch.Size([len(unique_index), len(inverse_mapping)]) features = spmm.apply(inverse_mapping, cols, vals, size, self._F) if (self.quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE): nums = spmm.apply( inverse_mapping, cols, vals, size, vals.reshape(N, 1), ) features /= nums elif self.quantization_mode == SparseTensorQuantizationMode.RANDOM_SUBSAMPLE: features = self._F[unique_index] else: # No quantization raise ValueError("Invalid quantization mode") sparse_tensor = SparseTensor( features, coordinate_map_key=sparse_tensor_key, coordinate_manager=self._manager, ) return sparse_tensor
def _get_coordinate_map_key(self, key_or_tensor_strides) -> CoordinateMapKey: r"""Helper function that retrieves the first coordinate map key for the given tensor stride.""" assert isinstance(key_or_tensor_strides, CoordinateMapKey) or isinstance( key_or_tensor_strides, (Sequence, np.ndarray, torch.IntTensor, int) ), f"The input must be either a CoordinateMapKey or tensor_stride of type (int, list, tuple, array, Tensor). Invalid: {key_or_tensor_strides}" if isinstance(key_or_tensor_strides, CoordinateMapKey): # Do nothing and return the input return key_or_tensor_strides else: tensor_strides = convert_to_int_list(key_or_tensor_strides, self.D) keys = self._manager.get_coordinate_map_keys(tensor_strides) assert len(keys) > 0 return keys[0]
def stride( self, coordinate_map_key: CoordinateMapKey, stride: Union[int, Sequence, np.ndarray, torch.Tensor], ) -> CoordinateMapKey: r"""Generate a new coordinate map and returns the key. :attr:`coordinate_map_key` (:attr:`MinkowskiEngine.CoordinateMapKey`): input map to generate the strided map from. :attr:`stride`: stride size. """ stride = convert_to_int_list(stride, self.D) return self._manager.stride(coordinate_map_key, stride)
def __init__( self, features: torch.Tensor, coordinates: torch.Tensor = None, # optional coordinate related arguments tensor_stride: StrideType = 1, coordinate_field_map_key: CoordinateMapKey = None, coordinate_manager: CoordinateManager = None, quantization_mode: SparseTensorQuantizationMode = SparseTensorQuantizationMode. UNWEIGHTED_AVERAGE, # optional manager related arguments allocator_type: GPUMemoryAllocatorType = None, minkowski_algorithm: MinkowskiAlgorithm = None, requires_grad=None, device=None, ): r""" Args: :attr:`features` (:attr:`torch.FloatTensor`, :attr:`torch.DoubleTensor`, :attr:`torch.cuda.FloatTensor`, or :attr:`torch.cuda.DoubleTensor`): The features of a sparse tensor. :attr:`coordinates` (:attr:`torch.IntTensor`): The coordinates associated to the features. If not provided, :attr:`coordinate_map_key` must be provided. :attr:`tensor_stride` (:attr:`int`, :attr:`list`, :attr:`numpy.array`, or :attr:`tensor.Tensor`): The tensor stride of the current sparse tensor. By default, it is 1. :attr:`coordinate_field_map_key` (:attr:`MinkowskiEngine.CoordinateMapKey`): When the coordinates are already cached in the MinkowskiEngine, we could reuse the same coordinate map by simply providing the coordinate map key. In most case, this process is done automatically. When you provide a `coordinate_field_map_key`, `coordinates` will be be ignored. :attr:`coordinate_manager` (:attr:`MinkowskiEngine.CoordinateManager`): The MinkowskiEngine manages all coordinate maps using the `_C.CoordinateMapManager`. If not provided, the MinkowskiEngine will create a new computation graph. In most cases, this process is handled automatically and you do not need to use this. :attr:`quantization_mode` (:attr:`MinkowskiEngine.SparseTensorQuantizationMode`): Defines how continuous coordinates will be quantized to define a sparse tensor. Please refer to :attr:`SparseTensorQuantizationMode` for details. :attr:`allocator_type` (:attr:`MinkowskiEngine.GPUMemoryAllocatorType`): Defines the GPU memory allocator type. By default, it uses the c10 allocator. :attr:`minkowski_algorithm` (:attr:`MinkowskiEngine.MinkowskiAlgorithm`): Controls the mode the minkowski engine runs, Use :attr:`MinkowskiAlgorithm.MEMORY_EFFICIENT` if you want to reduce the memory footprint. Or use :attr:`MinkowskiAlgorithm.SPEED_OPTIMIZED` if you want to make it run fasterat the cost of more memory. :attr:`requires_grad` (:attr:`bool`): Set the requires_grad flag. :attr:`device` (:attr:`torch.device`): Set the device the sparse tensor is defined. """ # Type checks assert isinstance(features, torch.Tensor), "Features must be a torch.Tensor" assert ( features.ndim == 2 ), f"The feature should be a matrix, The input feature is an order-{features.ndim} tensor." assert isinstance(quantization_mode, SparseTensorQuantizationMode) assert quantization_mode in [ SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE, SparseTensorQuantizationMode.UNWEIGHTED_SUM, SparseTensorQuantizationMode.RANDOM_SUBSAMPLE, SparseTensorQuantizationMode.MAX_POOL, ], "invalid quantization mode" self.quantization_mode = quantization_mode if coordinates is not None: assert isinstance(coordinates, torch.Tensor) if coordinate_field_map_key is not None: assert isinstance(coordinate_field_map_key, CoordinateMapKey) if coordinate_manager is not None: assert isinstance(coordinate_manager, CoordinateManager) if coordinates is None and (coordinate_field_map_key is None or coordinate_manager is None): raise ValueError( "Either coordinates or (coordinate_field_map_key, coordinate_manager) pair must be provided." ) Tensor.__init__(self) # To device if device is not None: features = features.to(device) if coordinates is not None: # assertion check for the map key done later coordinates = coordinates.to(device) self._D = (coordinates.size(1) - 1 if coordinates is not None else coordinate_manager.D) ########################## # Setup CoordsManager ########################## if coordinate_manager is None: # If set to share the coords man, use the global coords man if (sparse_tensor_operation_mode() == SparseTensorOperationMode.SHARE_COORDINATE_MANAGER): coordinate_manager = global_coordinate_manager() if coordinate_manager is None: coordinate_manager = CoordinateManager( D=self._D, coordinate_map_type=CoordinateMapType.CUDA if coordinates.is_cuda else CoordinateMapType.CPU, allocator_type=allocator_type, minkowski_algorithm=minkowski_algorithm, ) set_global_coordinate_manager(coordinate_manager) else: coordinate_manager = CoordinateManager( D=coordinates.size(1) - 1, coordinate_map_type=CoordinateMapType.CUDA if coordinates.is_cuda else CoordinateMapType.CPU, allocator_type=allocator_type, minkowski_algorithm=minkowski_algorithm, ) self._manager = coordinate_manager ########################## # Initialize coords ########################## # Coordinate Management if coordinates is not None: assert ( features.shape[0] == coordinates.shape[0] ), "The number of rows in features and coordinates must match." assert (features.is_cuda == coordinates.is_cuda ), "Features and coordinates must have the same backend." coordinate_field_map_key = CoordinateMapKey( convert_to_int_list(tensor_stride, self._D), "") coordinate_field_map_key = self._manager.insert_field( coordinates.float(), convert_to_int_list(tensor_stride, self._D), "") else: assert (coordinate_field_map_key.is_key_set() ), "The coordinate field map key must be valid." if requires_grad is not None: features.requires_grad_(requires_grad) self._F = features self._C = coordinates self.coordinate_field_map_key = coordinate_field_map_key self._batch_rows = None self._inverse_mapping = {}
def sparse( self, tensor_stride: Union[int, Sequence, np.array] = 1, coordinate_map_key: CoordinateMapKey = None, quantization_mode=None, ): r"""Converts the current sparse tensor field to a sparse tensor.""" if quantization_mode is None: quantization_mode = self.quantization_mode if coordinate_map_key is None: tensor_stride = convert_to_int_list(tensor_stride, self.D) coordinate_map_key, ( unique_index, inverse_mapping, ) = self._manager.field_to_sparse_insert_and_map( self.coordinate_field_map_key, tensor_stride, ) N_rows = len(unique_index) else: # sparse index, field index inverse_mapping, unique_index = self._manager.field_to_sparse_map( self.coordinate_field_map_key, coordinate_map_key, ) N_rows = self._manager.size(coordinate_map_key) self._inverse_mapping[coordinate_map_key] = inverse_mapping if quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_SUM: spmm = MinkowskiSPMMFunction() N = len(self._F) cols = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) vals = torch.ones(N, dtype=self._F.dtype, device=self._F.device) size = torch.Size([N_rows, len(inverse_mapping)]) features = spmm.apply(inverse_mapping, cols, vals, size, self._F) elif quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE: spmm_avg = MinkowskiSPMMAverageFunction() N = len(self._F) cols = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) size = torch.Size([N_rows, len(inverse_mapping)]) features = spmm_avg.apply(inverse_mapping, cols, size, self._F) elif quantization_mode == SparseTensorQuantizationMode.RANDOM_SUBSAMPLE: features = self._F[unique_index] elif quantization_mode == SparseTensorQuantizationMode.MAX_POOL: N = len(self._F) in_map = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) features = MinkowskiDirectMaxPoolingFunction().apply( in_map, inverse_mapping, self._F, N_rows) else: # No quantization raise ValueError("Invalid quantization mode") sparse_tensor = SparseTensor( features, coordinate_map_key=coordinate_map_key, coordinate_manager=self._manager, ) return sparse_tensor
def set_tensor_stride(self, s): ss = convert_to_int_list(s, self._D) self.coordinate_map_key.set_tensor_stride(ss)
def tensor_stride(self, p): r""" This function is not recommended to be used directly. """ p = convert_to_int_list(p, self._D) self.coordinate_map_key.set_tensor_stride(p)
def __init__( self, features: torch.Tensor, coordinates: torch.Tensor = None, # optional coordinate related arguments tensor_stride: StrideType = 1, coordinate_map_key: CoordinateMapKey = None, coordinate_manager: CoordinateManager = None, quantization_mode: SparseTensorQuantizationMode = SparseTensorQuantizationMode. RANDOM_SUBSAMPLE, # optional manager related arguments allocator_type: GPUMemoryAllocatorType = None, minkowski_algorithm: MinkowskiAlgorithm = None, device=None, ): r""" Args: :attr:`features` (:attr:`torch.FloatTensor`, :attr:`torch.DoubleTensor`, :attr:`torch.cuda.FloatTensor`, or :attr:`torch.cuda.DoubleTensor`): The features of a sparse tensor. :attr:`coordinates` (:attr:`torch.IntTensor`): The coordinates associated to the features. If not provided, :attr:`coordinate_map_key` must be provided. :attr:`coordinate_map_key` (:attr:`MinkowskiEngine.CoordinateMapKey`): When the coordinates are already cached in the MinkowskiEngine, we could reuse the same coordinate map by simply providing the coordinate map key. In most case, this process is done automatically. When you provide a `coordinate_map_key`, `coordinates` will be be ignored. :attr:`coordinate_manager` (:attr:`MinkowskiEngine.CoordinateManager`): The MinkowskiEngine manages all coordinate maps using the `_C.CoordinateMapManager`. If not provided, the MinkowskiEngine will create a new computation graph. In most cases, this process is handled automatically and you do not need to use this. :attr:`quantization_mode` (:attr:`MinkowskiEngine.SparseTensorQuantizationMode`): Defines how continuous coordinates will be quantized to define a sparse tensor. Please refer to :attr:`SparseTensorQuantizationMode` for details. :attr:`tensor_stride` (:attr:`int`, :attr:`list`, :attr:`numpy.array`, or :attr:`tensor.Tensor`): The tensor stride of the current sparse tensor. By default, it is 1. """ # Type checks assert isinstance(features, torch.Tensor), "Features must be a torch.Tensor" assert ( features.ndim == 2 ), f"The feature should be a matrix, The input feature is an order-{features.ndim} tensor." assert isinstance(quantization_mode, SparseTensorQuantizationMode) self.quantization_mode = quantization_mode if coordinates is not None: assert isinstance(coordinates, torch.Tensor) if coordinate_map_key is not None: assert isinstance(coordinate_map_key, CoordinateMapKey) if coordinate_manager is not None: assert isinstance(coordinate_manager, CoordinateManager) # To device if device is not None: features = features.to(device) coordinates = coordinates.to(device) # Coordinate Management self._D = 0 # coordinate size - 1 if coordinates is None and (coordinate_map_key is None or coordinate_manager is None): raise ValueError( "Either coordinates or (coordinate_map_key, coordinate_manager) pair must be provided." ) elif coordinates is not None: assert ( features.shape[0] == coordinates.shape[0] ), "The number of rows in features and coordinates must match." assert (features.is_cuda == coordinates.is_cuda ), "Features and coordinates must have the same backend." self._D = coordinates.size(1) - 1 coordinate_map_key = CoordinateMapKey( convert_to_int_list(tensor_stride, self._D), "") else: # not (coordinate_map_key is None or coordinate_manager is None) self._D = coordinate_manager.D ########################## # Setup CoordsManager ########################## if coordinate_manager is None: # If set to share the coords man, use the global coords man global _sparse_tensor_operation_mode, _global_coordinate_manager if (_sparse_tensor_operation_mode == SparseTensorOperationMode.SHARE_COORDINATE_MANAGER): if _global_coordinate_manager is None: _global_coordinate_manager = CoordinateManager( D=self._D, coordinate_map_type=CoordinateMapType.CUDA if coordinates.is_cuda else CoordinateMapType.CPU, allocator_type=allocator_type, minkowski_algorithm=minkowski_algorithm, ) coordinate_manager = _global_coordinate_manager else: coordinate_manager = CoordinateManager( D=coordinates.size(1) - 1, coordinate_map_type=CoordinateMapType.CUDA if coordinates.is_cuda else CoordinateMapType.CPU, allocator_type=allocator_type, minkowski_algorithm=minkowski_algorithm, ) self._manager = coordinate_manager ########################## # Initialize coords ########################## if coordinates is not None: coordinates, features, coordinate_map_key = self.initialize_coordinates( coordinates, features, coordinate_map_key) elif coordinate_map_key is not None: assert (coordinate_map_key.is_key_set() ), "The coordinate key must be a valid key." self.coordinate_map_key = coordinate_map_key self._F = features self._C = coordinates self._batch_rows = None
def sparse( self, tensor_stride: Union[int, Sequence, np.array] = 1, coordinate_map_key: CoordinateMapKey = None, quantization_mode: SparseTensorQuantizationMode = None, ): r"""Converts the current sparse tensor field to a sparse tensor.""" if quantization_mode is None: quantization_mode = self.quantization_mode assert ( quantization_mode != SparseTensorQuantizationMode.SPLAT_LINEAR_INTERPOLATION ), "Please use .splat() for splat quantization." if coordinate_map_key is None: tensor_stride = convert_to_int_list(tensor_stride, self.D) coordinate_map_key, ( unique_index, inverse_mapping, ) = self._manager.field_to_sparse_insert_and_map( self.coordinate_field_map_key, tensor_stride, ) N_rows = len(unique_index) else: # sparse index, field index inverse_mapping, unique_index = self._manager.field_to_sparse_map( self.coordinate_field_map_key, coordinate_map_key, ) N_rows = self._manager.size(coordinate_map_key) assert N_rows > 0, f"Invalid out coordinate map key. Found {N_row} elements." if len(inverse_mapping) == 0: # When the input has the same shape as the output self._inverse_mapping[coordinate_map_key] = torch.arange( len(self._F), dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) return SparseTensor( self._F, coordinate_map_key=coordinate_map_key, coordinate_manager=self._manager, ) # Create features if quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_SUM: N = len(self._F) cols = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) vals = torch.ones(N, dtype=self._F.dtype, device=self._F.device) size = torch.Size([N_rows, len(inverse_mapping)]) features = MinkowskiSPMMFunction().apply( inverse_mapping, cols, vals, size, self._F ) elif quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE: N = len(self._F) cols = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) size = torch.Size([N_rows, len(inverse_mapping)]) features = MinkowskiSPMMAverageFunction().apply( inverse_mapping, cols, size, self._F ) elif quantization_mode == SparseTensorQuantizationMode.RANDOM_SUBSAMPLE: features = self._F[unique_index] elif quantization_mode == SparseTensorQuantizationMode.MAX_POOL: N = len(self._F) in_map = torch.arange( N, dtype=inverse_mapping.dtype, device=inverse_mapping.device, ) features = MinkowskiDirectMaxPoolingFunction().apply( in_map, inverse_mapping, self._F, N_rows ) else: # No quantization raise ValueError("Invalid quantization mode") self._inverse_mapping[coordinate_map_key] = inverse_mapping return SparseTensor( features, coordinate_map_key=coordinate_map_key, coordinate_manager=self._manager, )