| from typing import * |
| import torch |
| import torch.nn as nn |
| from . import BACKEND, DEBUG |
| SparseTensorData = None |
|
|
|
|
| __all__ = [ |
| 'SparseTensor', |
| 'sparse_batch_broadcast', |
| 'sparse_batch_op', |
| 'sparse_cat', |
| 'sparse_unbind', |
| ] |
|
|
|
|
| class SparseTensor: |
| """ |
| Sparse tensor with support for both torchsparse and spconv backends. |
| |
| Parameters: |
| - feats (torch.Tensor): Features of the sparse tensor. |
| - coords (torch.Tensor): Coordinates of the sparse tensor. |
| - shape (torch.Size): Shape of the sparse tensor. |
| - layout (List[slice]): Layout of the sparse tensor for each batch |
| - data (SparseTensorData): Sparse tensor data used for convolusion |
| |
| NOTE: |
| - Data corresponding to a same batch should be contiguous. |
| - Coords should be in [0, 1023] |
| """ |
| @overload |
| def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ... |
|
|
| @overload |
| def __init__(self, data, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ... |
|
|
| def __init__(self, *args, **kwargs): |
| |
| global SparseTensorData |
| if SparseTensorData is None: |
| import importlib |
| if BACKEND == 'torchsparse': |
| SparseTensorData = importlib.import_module('torchsparse').SparseTensor |
| elif BACKEND == 'spconv': |
| SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor |
| |
| method_id = 0 |
| if len(args) != 0: |
| method_id = 0 if isinstance(args[0], torch.Tensor) else 1 |
| else: |
| method_id = 1 if 'data' in kwargs else 0 |
|
|
| if method_id == 0: |
| feats, coords, shape, layout = args + (None,) * (4 - len(args)) |
| if 'feats' in kwargs: |
| feats = kwargs['feats'] |
| del kwargs['feats'] |
| if 'coords' in kwargs: |
| coords = kwargs['coords'] |
| del kwargs['coords'] |
| if 'shape' in kwargs: |
| shape = kwargs['shape'] |
| del kwargs['shape'] |
| if 'layout' in kwargs: |
| layout = kwargs['layout'] |
| del kwargs['layout'] |
|
|
| if shape is None: |
| shape = self.__cal_shape(feats, coords) |
| if layout is None: |
| layout = self.__cal_layout(coords, shape[0]) |
| if BACKEND == 'torchsparse': |
| self.data = SparseTensorData(feats, coords, **kwargs) |
| elif BACKEND == 'spconv': |
| spatial_shape = list(coords.max(0)[0] + 1)[1:] |
| self.data = SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape, shape[0], **kwargs) |
| self.data._features = feats |
| elif method_id == 1: |
| data, shape, layout = args + (None,) * (3 - len(args)) |
| if 'data' in kwargs: |
| data = kwargs['data'] |
| del kwargs['data'] |
| if 'shape' in kwargs: |
| shape = kwargs['shape'] |
| del kwargs['shape'] |
| if 'layout' in kwargs: |
| layout = kwargs['layout'] |
| del kwargs['layout'] |
|
|
| self.data = data |
| if shape is None: |
| shape = self.__cal_shape(self.feats, self.coords) |
| if layout is None: |
| layout = self.__cal_layout(self.coords, shape[0]) |
|
|
| self._shape = shape |
| self._layout = layout |
| self._scale = kwargs.get('scale', (1, 1, 1)) |
| self._spatial_cache = kwargs.get('spatial_cache', {}) |
|
|
| if DEBUG: |
| try: |
| assert self.feats.shape[0] == self.coords.shape[0], f"Invalid feats shape: {self.feats.shape}, coords shape: {self.coords.shape}" |
| assert self.shape == self.__cal_shape(self.feats, self.coords), f"Invalid shape: {self.shape}" |
| assert self.layout == self.__cal_layout(self.coords, self.shape[0]), f"Invalid layout: {self.layout}" |
| for i in range(self.shape[0]): |
| assert torch.all(self.coords[self.layout[i], 0] == i), f"The data of batch {i} is not contiguous" |
| except Exception as e: |
| print('Debugging information:') |
| print(f"- Shape: {self.shape}") |
| print(f"- Layout: {self.layout}") |
| print(f"- Scale: {self._scale}") |
| print(f"- Coords: {self.coords}") |
| raise e |
| |
| def __cal_shape(self, feats, coords): |
| shape = [] |
| shape.append(coords[:, 0].max().item() + 1) |
| shape.extend([*feats.shape[1:]]) |
| return torch.Size(shape) |
| |
| def __cal_layout(self, coords, batch_size): |
| seq_len = torch.bincount(coords[:, 0], minlength=batch_size) |
| offset = torch.cumsum(seq_len, dim=0) |
| layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)] |
| return layout |
| |
| @property |
| def shape(self) -> torch.Size: |
| return self._shape |
| |
| def dim(self) -> int: |
| return len(self.shape) |
| |
| @property |
| def layout(self) -> List[slice]: |
| return self._layout |
|
|
| @property |
| def feats(self) -> torch.Tensor: |
| if BACKEND == 'torchsparse': |
| return self.data.F |
| elif BACKEND == 'spconv': |
| return self.data.features |
| |
| @feats.setter |
| def feats(self, value: torch.Tensor): |
| if BACKEND == 'torchsparse': |
| self.data.F = value |
| elif BACKEND == 'spconv': |
| self.data.features = value |
|
|
| @property |
| def coords(self) -> torch.Tensor: |
| if BACKEND == 'torchsparse': |
| return self.data.C |
| elif BACKEND == 'spconv': |
| return self.data.indices |
| |
| @coords.setter |
| def coords(self, value: torch.Tensor): |
| if BACKEND == 'torchsparse': |
| self.data.C = value |
| elif BACKEND == 'spconv': |
| self.data.indices = value |
|
|
| @property |
| def dtype(self): |
| return self.feats.dtype |
|
|
| @property |
| def device(self): |
| return self.feats.device |
|
|
| @overload |
| def to(self, dtype: torch.dtype) -> 'SparseTensor': ... |
|
|
| @overload |
| def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None) -> 'SparseTensor': ... |
|
|
| def to(self, *args, **kwargs) -> 'SparseTensor': |
| device = None |
| dtype = None |
| if len(args) == 2: |
| device, dtype = args |
| elif len(args) == 1: |
| if isinstance(args[0], torch.dtype): |
| dtype = args[0] |
| else: |
| device = args[0] |
| if 'dtype' in kwargs: |
| assert dtype is None, "to() received multiple values for argument 'dtype'" |
| dtype = kwargs['dtype'] |
| if 'device' in kwargs: |
| assert device is None, "to() received multiple values for argument 'device'" |
| device = kwargs['device'] |
| |
| new_feats = self.feats.to(device=device, dtype=dtype) |
| new_coords = self.coords.to(device=device) |
| return self.replace(new_feats, new_coords) |
|
|
| def type(self, dtype): |
| new_feats = self.feats.type(dtype) |
| return self.replace(new_feats) |
|
|
| def cpu(self) -> 'SparseTensor': |
| new_feats = self.feats.cpu() |
| new_coords = self.coords.cpu() |
| return self.replace(new_feats, new_coords) |
| |
| def cuda(self) -> 'SparseTensor': |
| new_feats = self.feats.cuda() |
| new_coords = self.coords.cuda() |
| return self.replace(new_feats, new_coords) |
|
|
| def half(self) -> 'SparseTensor': |
| new_feats = self.feats.half() |
| return self.replace(new_feats) |
| |
| def float(self) -> 'SparseTensor': |
| new_feats = self.feats.float() |
| return self.replace(new_feats) |
| |
| def detach(self) -> 'SparseTensor': |
| new_coords = self.coords.detach() |
| new_feats = self.feats.detach() |
| return self.replace(new_feats, new_coords) |
|
|
| def dense(self) -> torch.Tensor: |
| if BACKEND == 'torchsparse': |
| return self.data.dense() |
| elif BACKEND == 'spconv': |
| return self.data.dense() |
|
|
| def reshape(self, *shape) -> 'SparseTensor': |
| new_feats = self.feats.reshape(self.feats.shape[0], *shape) |
| return self.replace(new_feats) |
| |
| def unbind(self, dim: int) -> List['SparseTensor']: |
| return sparse_unbind(self, dim) |
|
|
| def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor': |
| new_shape = [self.shape[0]] |
| new_shape.extend(feats.shape[1:]) |
| if BACKEND == 'torchsparse': |
| new_data = SparseTensorData( |
| feats=feats, |
| coords=self.data.coords if coords is None else coords, |
| stride=self.data.stride, |
| spatial_range=self.data.spatial_range, |
| ) |
| new_data._caches = self.data._caches |
| elif BACKEND == 'spconv': |
| new_data = SparseTensorData( |
| self.data.features.reshape(self.data.features.shape[0], -1), |
| self.data.indices, |
| self.data.spatial_shape, |
| self.data.batch_size, |
| self.data.grid, |
| self.data.voxel_num, |
| self.data.indice_dict |
| ) |
| new_data._features = feats |
| new_data.benchmark = self.data.benchmark |
| new_data.benchmark_record = self.data.benchmark_record |
| new_data.thrust_allocator = self.data.thrust_allocator |
| new_data._timer = self.data._timer |
| new_data.force_algo = self.data.force_algo |
| new_data.int8_scale = self.data.int8_scale |
| if coords is not None: |
| new_data.indices = coords |
| new_tensor = SparseTensor(new_data, shape=torch.Size(new_shape), layout=self.layout, scale=self._scale, spatial_cache=self._spatial_cache) |
| return new_tensor |
|
|
| @staticmethod |
| def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor': |
| N, C = dim |
| x = torch.arange(aabb[0], aabb[3] + 1) |
| y = torch.arange(aabb[1], aabb[4] + 1) |
| z = torch.arange(aabb[2], aabb[5] + 1) |
| coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3) |
| coords = torch.cat([ |
| torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1), |
| coords.repeat(N, 1), |
| ], dim=1).to(dtype=torch.int32, device=device) |
| feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device) |
| return SparseTensor(feats=feats, coords=coords) |
|
|
| def __merge_sparse_cache(self, other: 'SparseTensor') -> dict: |
| new_cache = {} |
| for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())): |
| if k in self._spatial_cache: |
| new_cache[k] = self._spatial_cache[k] |
| if k in other._spatial_cache: |
| if k not in new_cache: |
| new_cache[k] = other._spatial_cache[k] |
| else: |
| new_cache[k].update(other._spatial_cache[k]) |
| return new_cache |
|
|
| def __neg__(self) -> 'SparseTensor': |
| return self.replace(-self.feats) |
| |
| def __elemwise__(self, other: Union[torch.Tensor, 'SparseTensor'], op: callable) -> 'SparseTensor': |
| if isinstance(other, torch.Tensor): |
| try: |
| other = torch.broadcast_to(other, self.shape) |
| other = sparse_batch_broadcast(self, other) |
| except: |
| pass |
| if isinstance(other, SparseTensor): |
| other = other.feats |
| new_feats = op(self.feats, other) |
| new_tensor = self.replace(new_feats) |
| if isinstance(other, SparseTensor): |
| new_tensor._spatial_cache = self.__merge_sparse_cache(other) |
| return new_tensor |
|
|
| def __add__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, torch.add) |
|
|
| def __radd__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, torch.add) |
| |
| def __sub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, torch.sub) |
| |
| def __rsub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, lambda x, y: torch.sub(y, x)) |
|
|
| def __mul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, torch.mul) |
|
|
| def __rmul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, torch.mul) |
|
|
| def __truediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, torch.div) |
|
|
| def __rtruediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor': |
| return self.__elemwise__(other, lambda x, y: torch.div(y, x)) |
|
|
| def __getitem__(self, idx): |
| if isinstance(idx, int): |
| idx = [idx] |
| elif isinstance(idx, slice): |
| idx = range(*idx.indices(self.shape[0])) |
| elif isinstance(idx, torch.Tensor): |
| if idx.dtype == torch.bool: |
| assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}" |
| idx = idx.nonzero().squeeze(1) |
| elif idx.dtype in [torch.int32, torch.int64]: |
| assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}" |
| else: |
| raise ValueError(f"Unknown index type: {idx.dtype}") |
| else: |
| raise ValueError(f"Unknown index type: {type(idx)}") |
| |
| coords = [] |
| feats = [] |
| for new_idx, old_idx in enumerate(idx): |
| coords.append(self.coords[self.layout[old_idx]].clone()) |
| coords[-1][:, 0] = new_idx |
| feats.append(self.feats[self.layout[old_idx]]) |
| coords = torch.cat(coords, dim=0).contiguous() |
| feats = torch.cat(feats, dim=0).contiguous() |
| return SparseTensor(feats=feats, coords=coords) |
|
|
| def register_spatial_cache(self, key, value) -> None: |
| """ |
| Register a spatial cache. |
| The spatial cache can be any thing you want to cache. |
| The registery and retrieval of the cache is based on current scale. |
| """ |
| scale_key = str(self._scale) |
| if scale_key not in self._spatial_cache: |
| self._spatial_cache[scale_key] = {} |
| self._spatial_cache[scale_key][key] = value |
|
|
| def get_spatial_cache(self, key=None): |
| """ |
| Get a spatial cache. |
| """ |
| scale_key = str(self._scale) |
| cur_scale_cache = self._spatial_cache.get(scale_key, {}) |
| if key is None: |
| return cur_scale_cache |
| return cur_scale_cache.get(key, None) |
|
|
|
|
| def sparse_batch_broadcast(input: SparseTensor, other: torch.Tensor) -> torch.Tensor: |
| """ |
| Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation. |
| |
| Args: |
| input (torch.Tensor): 1D tensor to broadcast. |
| target (SparseTensor): Sparse tensor to broadcast to. |
| op (callable): Operation to perform after broadcasting. Defaults to torch.add. |
| """ |
| coords, feats = input.coords, input.feats |
| broadcasted = torch.zeros_like(feats) |
| for k in range(input.shape[0]): |
| broadcasted[input.layout[k]] = other[k] |
| return broadcasted |
|
|
|
|
| def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor: |
| """ |
| Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation. |
| |
| Args: |
| input (torch.Tensor): 1D tensor to broadcast. |
| target (SparseTensor): Sparse tensor to broadcast to. |
| op (callable): Operation to perform after broadcasting. Defaults to torch.add. |
| """ |
| return input.replace(op(input.feats, sparse_batch_broadcast(input, other))) |
|
|
|
|
| def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor: |
| """ |
| Concatenate a list of sparse tensors. |
| |
| Args: |
| inputs (List[SparseTensor]): List of sparse tensors to concatenate. |
| """ |
| if dim == 0: |
| start = 0 |
| coords = [] |
| for input in inputs: |
| coords.append(input.coords.clone()) |
| coords[-1][:, 0] += start |
| start += input.shape[0] |
| coords = torch.cat(coords, dim=0) |
| feats = torch.cat([input.feats for input in inputs], dim=0) |
| output = SparseTensor( |
| coords=coords, |
| feats=feats, |
| ) |
| else: |
| feats = torch.cat([input.feats for input in inputs], dim=dim) |
| output = inputs[0].replace(feats) |
|
|
| return output |
|
|
|
|
| def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]: |
| """ |
| Unbind a sparse tensor along a dimension. |
| |
| Args: |
| input (SparseTensor): Sparse tensor to unbind. |
| dim (int): Dimension to unbind. |
| """ |
| if dim == 0: |
| return [input[i] for i in range(input.shape[0])] |
| else: |
| feats = input.feats.unbind(dim) |
| return [input.replace(f) for f in feats] |
|
|
