| | from typing import * |
| | from fractions import Fraction |
| | import torch |
| | from . import config |
| |
|
| |
|
| | __all__ = [ |
| | 'VarLenTensor', |
| | 'varlen_cat', |
| | 'varlen_unbind', |
| | 'SparseTensor', |
| | 'sparse_cat', |
| | 'sparse_unbind', |
| | ] |
| |
|
| |
|
| | class VarLenTensor: |
| | """ |
| | Sequential tensor with variable length. |
| | |
| | Args: |
| | feats (torch.Tensor): Features of the varlen tensor. |
| | layout (List[slice]): Layout of the varlen tensor for each batch |
| | """ |
| | def __init__(self, feats: torch.Tensor, layout: List[slice]=None): |
| | self.feats = feats |
| | self.layout = layout if layout is not None else [slice(0, feats.shape[0])] |
| | self._cache = {} |
| | |
| | @staticmethod |
| | def layout_from_seqlen(seqlen: list) -> List[slice]: |
| | """ |
| | Create a layout from a tensor of sequence lengths. |
| | """ |
| | layout = [] |
| | start = 0 |
| | for l in seqlen: |
| | layout.append(slice(start, start + l)) |
| | start += l |
| | return layout |
| | |
| | @staticmethod |
| | def from_tensor_list(tensor_list: List[torch.Tensor]) -> 'VarLenTensor': |
| | """ |
| | Create a VarLenTensor from a list of tensors. |
| | """ |
| | feats = torch.cat(tensor_list, dim=0) |
| | layout = [] |
| | start = 0 |
| | for tensor in tensor_list: |
| | layout.append(slice(start, start + tensor.shape[0])) |
| | start += tensor.shape[0] |
| | return VarLenTensor(feats, layout) |
| | |
| | def to_tensor_list(self) -> List[torch.Tensor]: |
| | """ |
| | Convert a VarLenTensor to a list of tensors. |
| | """ |
| | tensor_list = [] |
| | for s in self.layout: |
| | tensor_list.append(self.feats[s]) |
| | return tensor_list |
| | |
| | def __len__(self) -> int: |
| | return len(self.layout) |
| | |
| | @property |
| | def shape(self) -> torch.Size: |
| | return torch.Size([len(self.layout), *self.feats.shape[1:]]) |
| | |
| | def dim(self) -> int: |
| | return len(self.shape) |
| | |
| | @property |
| | def ndim(self) -> int: |
| | return self.dim() |
| |
|
| | @property |
| | def dtype(self): |
| | return self.feats.dtype |
| |
|
| | @property |
| | def device(self): |
| | return self.feats.device |
| | |
| | @property |
| | def seqlen(self) -> torch.LongTensor: |
| | if 'seqlen' not in self._cache: |
| | self._cache['seqlen'] = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device) |
| | return self._cache['seqlen'] |
| | |
| | @property |
| | def cum_seqlen(self) -> torch.LongTensor: |
| | if 'cum_seqlen' not in self._cache: |
| | self._cache['cum_seqlen'] = torch.cat([ |
| | torch.tensor([0], dtype=torch.long, device=self.device), |
| | self.seqlen.cumsum(dim=0) |
| | ], dim=0) |
| | return self._cache['cum_seqlen'] |
| | |
| | @property |
| | def batch_boardcast_map(self) -> torch.LongTensor: |
| | """ |
| | Get the broadcast map for the varlen tensor. |
| | """ |
| | if 'batch_boardcast_map' not in self._cache: |
| | self._cache['batch_boardcast_map'] = torch.repeat_interleave( |
| | torch.arange(len(self.layout), device=self.device), |
| | self.seqlen, |
| | ) |
| | return self._cache['batch_boardcast_map'] |
| | |
| | @overload |
| | def to(self, dtype: torch.dtype, *, non_blocking: bool = False, copy: bool = False) -> 'VarLenTensor': ... |
| |
|
| | @overload |
| | def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None, *, non_blocking: bool = False, copy: bool = False) -> 'VarLenTensor': ... |
| |
|
| | def to(self, *args, **kwargs) -> 'VarLenTensor': |
| | 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'] |
| | non_blocking = kwargs.get('non_blocking', False) |
| | copy = kwargs.get('copy', False) |
| | |
| | new_feats = self.feats.to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy) |
| | return self.replace(new_feats) |
| |
|
| | def type(self, dtype): |
| | new_feats = self.feats.type(dtype) |
| | return self.replace(new_feats) |
| |
|
| | def cpu(self) -> 'VarLenTensor': |
| | new_feats = self.feats.cpu() |
| | return self.replace(new_feats) |
| | |
| | def cuda(self) -> 'VarLenTensor': |
| | new_feats = self.feats.cuda() |
| | return self.replace(new_feats) |
| |
|
| | def half(self) -> 'VarLenTensor': |
| | new_feats = self.feats.half() |
| | return self.replace(new_feats) |
| | |
| | def float(self) -> 'VarLenTensor': |
| | new_feats = self.feats.float() |
| | return self.replace(new_feats) |
| | |
| | def detach(self) -> 'VarLenTensor': |
| | new_feats = self.feats.detach() |
| | return self.replace(new_feats) |
| |
|
| | def reshape(self, *shape) -> 'VarLenTensor': |
| | new_feats = self.feats.reshape(self.feats.shape[0], *shape) |
| | return self.replace(new_feats) |
| | |
| | def unbind(self, dim: int) -> List['VarLenTensor']: |
| | return varlen_unbind(self, dim) |
| |
|
| | def replace(self, feats: torch.Tensor) -> 'VarLenTensor': |
| | new_tensor = VarLenTensor( |
| | feats=feats, |
| | layout=self.layout, |
| | ) |
| | new_tensor._cache = self._cache |
| | return new_tensor |
| | |
| | def to_dense(self, max_length=None) -> torch.Tensor: |
| | """ |
| | Convert a VarLenTensor to a dense representation without for-loop. |
| | |
| | Returns: |
| | dense (torch.Tensor): (N, L, C) dense tensor |
| | mask (torch.BoolTensor): (N, L) mask indicating valid positions |
| | """ |
| | N = len(self) |
| | L = max_length or self.seqlen.max().item() |
| | spatial = self.feats.shape[1:] |
| | idx = torch.arange(L, device=self.device).unsqueeze(0).expand(N, L) |
| | mask = (idx < self.seqlen.unsqueeze(1)) |
| | mapping = mask.reshape(-1).cumsum(dim=0) - 1 |
| | dense = self.feats[mapping] |
| | dense = dense.reshape(N, L, *spatial) |
| | return dense, mask |
| |
|
| | def __neg__(self) -> 'VarLenTensor': |
| | return self.replace(-self.feats) |
| | |
| | def __elemwise__(self, other: Union[torch.Tensor, 'VarLenTensor'], op: callable) -> 'VarLenTensor': |
| | if isinstance(other, torch.Tensor): |
| | try: |
| | other = torch.broadcast_to(other, self.shape) |
| | other = other[self.batch_boardcast_map] |
| | except: |
| | pass |
| | if isinstance(other, VarLenTensor): |
| | other = other.feats |
| | new_feats = op(self.feats, other) |
| | new_tensor = self.replace(new_feats) |
| | return new_tensor |
| |
|
| | def __add__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, torch.add) |
| |
|
| | def __radd__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, torch.add) |
| | |
| | def __sub__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, torch.sub) |
| | |
| | def __rsub__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, lambda x, y: torch.sub(y, x)) |
| |
|
| | def __mul__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, torch.mul) |
| |
|
| | def __rmul__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, torch.mul) |
| |
|
| | def __truediv__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | return self.__elemwise__(other, torch.div) |
| |
|
| | def __rtruediv__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor': |
| | 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, list): |
| | assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}" |
| | 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)}") |
| | |
| | new_feats = [] |
| | new_layout = [] |
| | start = 0 |
| | for new_idx, old_idx in enumerate(idx): |
| | new_feats.append(self.feats[self.layout[old_idx]]) |
| | new_layout.append(slice(start, start + len(new_feats[-1]))) |
| | start += len(new_feats[-1]) |
| | new_feats = torch.cat(new_feats, dim=0).contiguous() |
| | new_tensor = VarLenTensor(feats=new_feats, layout=new_layout) |
| | return new_tensor |
| | |
| | def reduce(self, op: str, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: |
| | if isinstance(dim, int): |
| | dim = (dim,) |
| | |
| | if op =='mean': |
| | red = self.feats.mean(dim=dim, keepdim=keepdim) |
| | elif op =='sum': |
| | red = self.feats.sum(dim=dim, keepdim=keepdim) |
| | elif op == 'prod': |
| | red = self.feats.prod(dim=dim, keepdim=keepdim) |
| | else: |
| | raise ValueError(f"Unsupported reduce operation: {op}") |
| | |
| | if dim is None or 0 in dim: |
| | return red |
| | |
| | red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen) |
| | return red |
| | |
| | def mean(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: |
| | return self.reduce(op='mean', dim=dim, keepdim=keepdim) |
| | |
| | def sum(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: |
| | return self.reduce(op='sum', dim=dim, keepdim=keepdim) |
| | |
| | def prod(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: |
| | return self.reduce(op='prod', dim=dim, keepdim=keepdim) |
| | |
| | def std(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor: |
| | mean = self.mean(dim=dim, keepdim=True) |
| | mean2 = self.replace(self.feats ** 2).mean(dim=dim, keepdim=True) |
| | std = (mean2 - mean ** 2).sqrt() |
| | return std |
| | |
| | def __repr__(self) -> str: |
| | return f"VarLenTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})" |
| |
|
| |
|
| | def varlen_cat(inputs: List[VarLenTensor], dim: int = 0) -> VarLenTensor: |
| | """ |
| | Concatenate a list of varlen tensors. |
| | |
| | Args: |
| | inputs (List[VarLenTensor]): List of varlen tensors to concatenate. |
| | """ |
| | if dim == 0: |
| | new_feats = torch.cat([input.feats for input in inputs], dim=0) |
| | start = 0 |
| | new_layout = [] |
| | for input in inputs: |
| | for l in input.layout: |
| | new_layout.append(slice(start, start + l.stop - l.start)) |
| | start += l.stop - l.start |
| | output = VarLenTensor(feats=new_feats, layout=new_layout) |
| | else: |
| | feats = torch.cat([input.feats for input in inputs], dim=dim) |
| | output = inputs[0].replace(feats) |
| |
|
| | return output |
| |
|
| |
|
| | def varlen_unbind(input: VarLenTensor, dim: int) -> Union[List[VarLenTensor]]: |
| | """ |
| | Unbind a varlen tensor along a dimension. |
| | |
| | Args: |
| | input (VarLenTensor): Varlen tensor to unbind. |
| | dim (int): Dimension to unbind. |
| | """ |
| | if dim == 0: |
| | return [input[i] for i in range(len(input))] |
| | else: |
| | feats = input.feats.unbind(dim) |
| | return [input.replace(f) for f in feats] |
| | |
| |
|
| | class SparseTensor(VarLenTensor): |
| | """ |
| | 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] |
| | """ |
| | SparseTensorData = None |
| |
|
| | @overload |
| | def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, **kwargs): ... |
| |
|
| | @overload |
| | def __init__(self, data, shape: Optional[torch.Size] = None, **kwargs): ... |
| |
|
| | def __init__(self, *args, **kwargs): |
| | |
| | if self.SparseTensorData is None: |
| | import importlib |
| | if config.CONV == 'torchsparse': |
| | self.SparseTensorData = importlib.import_module('torchsparse').SparseTensor |
| | elif config.CONV == 'spconv': |
| | self.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 = args + (None,) * (3 - 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 config.CONV == 'torchsparse': |
| | self.data = self.SparseTensorData(feats, coords, **kwargs) |
| | elif config.CONV == 'spconv': |
| | spatial_shape = list(coords.max(0)[0] + 1) |
| | self.data = self.SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape[1:], spatial_shape[0], **kwargs) |
| | self.data._features = feats |
| | else: |
| | self.data = { |
| | 'feats': feats, |
| | 'coords': coords, |
| | } |
| | elif method_id == 1: |
| | data, shape = args + (None,) * (2 - len(args)) |
| | if 'data' in kwargs: |
| | data = kwargs['data'] |
| | del kwargs['data'] |
| | if 'shape' in kwargs: |
| | shape = kwargs['shape'] |
| | del kwargs['shape'] |
| |
|
| | self.data = data |
| |
|
| | self._shape = shape |
| | self._scale = kwargs.get('scale', (Fraction(1, 1), Fraction(1, 1), Fraction(1, 1))) |
| | self._spatial_cache = kwargs.get('spatial_cache', {}) |
| |
|
| | if config.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 |
| | |
| | @staticmethod |
| | def from_tensor_list(feats_list: List[torch.Tensor], coords_list: List[torch.Tensor]) -> 'SparseTensor': |
| | """ |
| | Create a SparseTensor from a list of tensors. |
| | """ |
| | feats = torch.cat(feats_list, dim=0) |
| | coords = [] |
| | for i, coord in enumerate(coords_list): |
| | coord = torch.cat([torch.full_like(coord[:, :1], i), coord[:, 1:]], dim=1) |
| | coords.append(coord) |
| | coords = torch.cat(coords, dim=0) |
| | return SparseTensor(feats, coords) |
| | |
| | def to_tensor_list(self) -> Tuple[List[torch.Tensor], List[torch.Tensor]]: |
| | """ |
| | Convert a SparseTensor to list of tensors. |
| | """ |
| | feats_list = [] |
| | coords_list = [] |
| | for s in self.layout: |
| | feats_list.append(self.feats[s]) |
| | coords_list.append(self.coords[s]) |
| | return feats_list, coords_list |
| | |
| | def __len__(self) -> int: |
| | return len(self.layout) |
| | |
| | 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 |
| | |
| | def __cal_spatial_shape(self, coords): |
| | return torch.Size((coords[:, 1:].max(0)[0] + 1).tolist()) |
| | |
| | @property |
| | def shape(self) -> torch.Size: |
| | if self._shape is None: |
| | self._shape = self.__cal_shape(self.feats, self.coords) |
| | return self._shape |
| | |
| | @property |
| | def layout(self) -> List[slice]: |
| | layout = self.get_spatial_cache('layout') |
| | if layout is None: |
| | layout = self.__cal_layout(self.coords, self.shape[0]) |
| | self.register_spatial_cache('layout', layout) |
| | return layout |
| | |
| | @property |
| | def spatial_shape(self) -> torch.Size: |
| | spatial_shape = self.get_spatial_cache('shape') |
| | if spatial_shape is None: |
| | spatial_shape = self.__cal_spatial_shape(self.coords) |
| | self.register_spatial_cache('shape', spatial_shape) |
| | return spatial_shape |
| |
|
| | @property |
| | def feats(self) -> torch.Tensor: |
| | if config.CONV == 'torchsparse': |
| | return self.data.F |
| | elif config.CONV == 'spconv': |
| | return self.data.features |
| | else: |
| | return self.data['feats'] |
| | |
| | @feats.setter |
| | def feats(self, value: torch.Tensor): |
| | if config.CONV == 'torchsparse': |
| | self.data.F = value |
| | elif config.CONV == 'spconv': |
| | self.data.features = value |
| | else: |
| | self.data['feats'] = value |
| |
|
| | @property |
| | def coords(self) -> torch.Tensor: |
| | if config.CONV == 'torchsparse': |
| | return self.data.C |
| | elif config.CONV == 'spconv': |
| | return self.data.indices |
| | else: |
| | return self.data['coords'] |
| | |
| | @coords.setter |
| | def coords(self, value: torch.Tensor): |
| | if config.CONV == 'torchsparse': |
| | self.data.C = value |
| | elif config.CONV == 'spconv': |
| | self.data.indices = value |
| | else: |
| | self.data['coords'] = value |
| |
|
| | @property |
| | def dtype(self): |
| | return self.feats.dtype |
| |
|
| | @property |
| | def device(self): |
| | return self.feats.device |
| | |
| | @property |
| | def seqlen(self) -> torch.LongTensor: |
| | seqlen = self.get_spatial_cache('seqlen') |
| | if seqlen is None: |
| | seqlen = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device) |
| | self.register_spatial_cache('seqlen', seqlen) |
| | return seqlen |
| | |
| | @property |
| | def cum_seqlen(self) -> torch.LongTensor: |
| | cum_seqlen = self.get_spatial_cache('cum_seqlen') |
| | if cum_seqlen is None: |
| | cum_seqlen = torch.cat([ |
| | torch.tensor([0], dtype=torch.long, device=self.device), |
| | self.seqlen.cumsum(dim=0) |
| | ], dim=0) |
| | self.register_spatial_cache('cum_seqlen', cum_seqlen) |
| | return cum_seqlen |
| | |
| | @property |
| | def batch_boardcast_map(self) -> torch.LongTensor: |
| | """ |
| | Get the broadcast map for the varlen tensor. |
| | """ |
| | batch_boardcast_map = self.get_spatial_cache('batch_boardcast_map') |
| | if batch_boardcast_map is None: |
| | batch_boardcast_map = torch.repeat_interleave( |
| | torch.arange(len(self.layout), device=self.device), |
| | self.seqlen, |
| | ) |
| | self.register_spatial_cache('batch_boardcast_map', batch_boardcast_map) |
| | return batch_boardcast_map |
| |
|
| | @overload |
| | def to(self, dtype: torch.dtype, *, non_blocking: bool = False, copy: bool = False) -> 'SparseTensor': ... |
| |
|
| | @overload |
| | def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None, *, non_blocking: bool = False, copy: bool = False) -> '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'] |
| | non_blocking = kwargs.get('non_blocking', False) |
| | copy = kwargs.get('copy', False) |
| | |
| | new_feats = self.feats.to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy) |
| | new_coords = self.coords.to(device=device, non_blocking=non_blocking, copy=copy) |
| | 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 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': |
| | if config.CONV == 'torchsparse': |
| | new_data = self.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 config.CONV == 'spconv': |
| | new_data = self.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 |
| | else: |
| | new_data = { |
| | 'feats': feats, |
| | 'coords': self.data['coords'] if coords is None else coords, |
| | } |
| | new_tensor = SparseTensor( |
| | new_data, |
| | shape=torch.Size([self._shape[0]] + list(feats.shape[1:])) if self._shape is not None else None, |
| | scale=self._scale, |
| | spatial_cache=self._spatial_cache |
| | ) |
| | return new_tensor |
| | |
| | def to_dense(self) -> torch.Tensor: |
| | if config.CONV == 'torchsparse': |
| | return self.data.dense() |
| | elif config.CONV == 'spconv': |
| | return self.data.dense() |
| | else: |
| | spatial_shape = self.spatial_shape |
| | ret = torch.zeros(*self.shape, *spatial_shape, dtype=self.dtype, device=self.device) |
| | idx = [self.coords[:, 0], slice(None)] + self.coords[:, 1:].unbind(1) |
| | ret[tuple(idx)] = self.feats |
| | return ret |
| |
|
| | @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 __elemwise__(self, other: Union[torch.Tensor, VarLenTensor], op: callable) -> 'SparseTensor': |
| | if isinstance(other, torch.Tensor): |
| | try: |
| | other = torch.broadcast_to(other, self.shape) |
| | other = other[self.batch_boardcast_map] |
| | except: |
| | pass |
| | if isinstance(other, VarLenTensor): |
| | 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 __getitem__(self, idx): |
| | if isinstance(idx, int): |
| | idx = [idx] |
| | elif isinstance(idx, slice): |
| | idx = range(*idx.indices(self.shape[0])) |
| | elif isinstance(idx, list): |
| | assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}" |
| | 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)}") |
| | |
| | new_coords = [] |
| | new_feats = [] |
| | new_layout = [] |
| | new_shape = torch.Size([len(idx)] + list(self.shape[1:])) |
| | start = 0 |
| | for new_idx, old_idx in enumerate(idx): |
| | new_coords.append(self.coords[self.layout[old_idx]].clone()) |
| | new_coords[-1][:, 0] = new_idx |
| | new_feats.append(self.feats[self.layout[old_idx]]) |
| | new_layout.append(slice(start, start + len(new_coords[-1]))) |
| | start += len(new_coords[-1]) |
| | new_coords = torch.cat(new_coords, dim=0).contiguous() |
| | new_feats = torch.cat(new_feats, dim=0).contiguous() |
| | new_tensor = SparseTensor(feats=new_feats, coords=new_coords, shape=new_shape) |
| | new_tensor.register_spatial_cache('layout', new_layout) |
| | return new_tensor |
| | |
| | def clear_spatial_cache(self) -> None: |
| | """ |
| | Clear all spatial caches. |
| | """ |
| | self._spatial_cache = {} |
| |
|
| | 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 __repr__(self) -> str: |
| | return f"SparseTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})" |
| |
|
| | 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] |
| |
|
