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]): L
| 15 | |
| 16 | |
| 17 | class SparseTensor: |
| 18 | """ |
| 19 | Sparse tensor with support for both torchsparse and spconv backends. |
| 20 | |
| 21 | Parameters: |
| 22 | - feats (torch.Tensor): Features of the sparse tensor. |
| 23 | - coords (torch.Tensor): Coordinates of the sparse tensor. |
| 24 | - shape (torch.Size): Shape of the sparse tensor. |
| 25 | - layout (List[slice]): Layout of the sparse tensor for each batch |
| 26 | - data (SparseTensorData): Sparse tensor data used for convolusion |
| 27 | |
| 28 | NOTE: |
| 29 | - Data corresponding to a same batch should be contiguous. |
| 30 | - Coords should be in [0, 1023] |
| 31 | """ |
| 32 | @overload |
| 33 | def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ... |
| 34 | |
| 35 | @overload |
| 36 | def __init__(self, data, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ... |
| 37 | |
| 38 | def __init__(self, *args, **kwargs): |
| 39 | # Lazy import of sparse tensor backend |
| 40 | global SparseTensorData |
| 41 | if SparseTensorData is None: |
| 42 | import importlib |
| 43 | if BACKEND == 'torchsparse': |
| 44 | SparseTensorData = importlib.import_module('torchsparse').SparseTensor |
| 45 | elif BACKEND == 'spconv': |
| 46 | SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor |
| 47 | |
| 48 | method_id = 0 |
| 49 | if len(args) != 0: |
| 50 | method_id = 0 if isinstance(args[0], torch.Tensor) else 1 |
| 51 | else: |
| 52 | method_id = 1 if 'data' in kwargs else 0 |
| 53 | |
| 54 | if method_id == 0: |
| 55 | feats, coords, shape, layout = args + (None,) * (4 - len(args)) |
| 56 | if 'feats' in kwargs: |
| 57 | feats = kwargs['feats'] |
| 58 | del kwargs['feats'] |
| 59 | if 'coords' in kwargs: |
| 60 | coords = kwargs['coords'] |
| 61 | del kwargs['coords'] |
| 62 | if 'shape' in kwargs: |
| 63 | shape = kwargs['shape'] |
| 64 | del kwargs['shape'] |
| 65 | if 'layout' in kwargs: |
| 66 | layout = kwargs['layout'] |
| 67 | del kwargs['layout'] |
| 68 | |
| 69 | if shape is None: |
| 70 | shape = self.__cal_shape(feats, coords) |
| 71 | if layout is None: |
| 72 | layout = self.__cal_layout(coords, shape[0]) |
| 73 | if BACKEND == 'torchsparse': |
| 74 | self.data = SparseTensorData(feats, coords, **kwargs) |
no outgoing calls
no test coverage detected