1383 lines
54 KiB
ReStructuredText
1383 lines
54 KiB
ReStructuredText
.. automodule:: torch.sparse
|
|
|
|
.. currentmodule:: torch
|
|
|
|
.. _sparse-docs:
|
|
|
|
torch.sparse
|
|
============
|
|
|
|
.. warning::
|
|
|
|
The PyTorch API of sparse tensors is in beta and may change in the near future.
|
|
We highly welcome feature requests, bug reports and general suggestions as GitHub issues.
|
|
|
|
Why and when to use sparsity
|
|
++++++++++++++++++++++++++++
|
|
|
|
By default PyTorch stores :class:`torch.Tensor` stores elements contiguously
|
|
physical memory. This leads to efficient implementations of various array
|
|
processing algorithms that require fast access to elements.
|
|
|
|
Now, some users might decide to represent data such as graph adjacency
|
|
matrices, pruned weights or points clouds by Tensors whose *elements are
|
|
mostly zero valued*. We recognize these are important applications and aim
|
|
to provide performance optimizations for these use cases via sparse storage formats.
|
|
|
|
Various sparse storage formats such as COO, CSR/CSC, semi-structured, LIL, etc. have been
|
|
developed over the years. While they differ in exact layouts, they all
|
|
compress data through efficient representation of zero valued elements.
|
|
We call the uncompressed values *specified* in contrast to *unspecified*,
|
|
compressed elements.
|
|
|
|
By compressing repeat zeros sparse storage formats aim to save memory
|
|
and computational resources on various CPUs and GPUs. Especially for high
|
|
degrees of sparsity or highly structured sparsity this can have significant
|
|
performance implications. As such sparse storage formats can be seen as a
|
|
performance optimization.
|
|
|
|
Like many other performance optimization sparse storage formats are not
|
|
always advantageous. When trying sparse formats for your use case
|
|
you might find your execution time to increase rather than decrease.
|
|
|
|
Please feel encouraged to open a GitHub issue if you analytically
|
|
expected to see a stark increase in performance but measured a
|
|
degradation instead. This helps us prioritize the implementation
|
|
of efficient kernels and wider performance optimizations.
|
|
|
|
We make it easy to try different sparsity layouts, and convert between them,
|
|
without being opinionated on what's best for your particular application.
|
|
|
|
Functionality overview
|
|
++++++++++++++++++++++
|
|
|
|
We want it to be straightforward to construct a sparse Tensor from a
|
|
given dense Tensor by providing conversion routines for each layout.
|
|
|
|
In the next example we convert a 2D Tensor with default dense (strided)
|
|
layout to a 2D Tensor backed by the COO memory layout. Only values and
|
|
indices of non-zero elements are stored in this case.
|
|
|
|
>>> a = torch.tensor([[0, 2.], [3, 0]])
|
|
>>> a.to_sparse()
|
|
tensor(indices=tensor([[0, 1],
|
|
[1, 0]]),
|
|
values=tensor([2., 3.]),
|
|
size=(2, 2), nnz=2, layout=torch.sparse_coo)
|
|
|
|
PyTorch currently supports :ref:`COO<sparse-coo-docs>`, :ref:`CSR<sparse-csr-docs>`,
|
|
:ref:`CSC<sparse-csc-docs>`, :ref:`BSR<sparse-bsr-docs>`, and :ref:`BSC<sparse-bsc-docs>`.
|
|
|
|
We also have a prototype implementation to support :ref: `semi-structured sparsity<sparse-semi-structured-docs>`.
|
|
Please see the references for more details.
|
|
|
|
Note that we provide slight generalizations of these formats.
|
|
|
|
Batching: Devices such as GPUs require batching for optimal performance and
|
|
thus we support batch dimensions.
|
|
|
|
We currently offer a very simple version of batching where each component of a sparse format
|
|
itself is batched. This also requires the same number of specified elements per batch entry.
|
|
In this example we construct a 3D (batched) CSR Tensor from a 3D dense Tensor.
|
|
|
|
>>> t = torch.tensor([[[1., 0], [2., 3.]], [[4., 0], [5., 6.]]])
|
|
>>> t.dim()
|
|
3
|
|
>>> t.to_sparse_csr()
|
|
tensor(crow_indices=tensor([[0, 1, 3],
|
|
[0, 1, 3]]),
|
|
col_indices=tensor([[0, 0, 1],
|
|
[0, 0, 1]]),
|
|
values=tensor([[1., 2., 3.],
|
|
[4., 5., 6.]]), size=(2, 2, 2), nnz=3,
|
|
layout=torch.sparse_csr)
|
|
|
|
|
|
Dense dimensions: On the other hand, some data such as Graph embeddings might be
|
|
better viewed as sparse collections of vectors instead of scalars.
|
|
|
|
In this example we create a 3D Hybrid COO Tensor with 2 sparse and 1 dense dimension
|
|
from a 3D strided Tensor. If an entire row in the 3D strided Tensor is zero, it is
|
|
not stored. If however any of the values in the row are non-zero, they are stored
|
|
entirely. This reduces the number of indices since we need one index one per row instead
|
|
of one per element. But it also increases the amount of storage for the values. Since
|
|
only rows that are *entirely* zero can be emitted and the presence of any non-zero
|
|
valued elements cause the entire row to be stored.
|
|
|
|
>>> t = torch.tensor([[[0., 0], [1., 2.]], [[0., 0], [3., 4.]]])
|
|
>>> t.to_sparse(sparse_dim=2)
|
|
tensor(indices=tensor([[0, 1],
|
|
[1, 1]]),
|
|
values=tensor([[1., 2.],
|
|
[3., 4.]]),
|
|
size=(2, 2, 2), nnz=2, layout=torch.sparse_coo)
|
|
|
|
|
|
Operator overview
|
|
+++++++++++++++++
|
|
|
|
Fundamentally, operations on Tensor with sparse storage formats behave the same as
|
|
operations on Tensor with strided (or other) storage formats. The particularities of
|
|
storage, that is the physical layout of the data, influences the performance of
|
|
an operation but should not influence the semantics.
|
|
|
|
|
|
We are actively increasing operator coverage for sparse tensors. Users should not
|
|
expect support same level of support as for dense Tensors yet.
|
|
See our :ref:`operator<sparse-ops-docs>` documentation for a list.
|
|
|
|
>>> b = torch.tensor([[0, 0, 1, 2, 3, 0], [4, 5, 0, 6, 0, 0]])
|
|
>>> b_s = b.to_sparse_csr()
|
|
>>> b_s.cos()
|
|
Traceback (most recent call last):
|
|
File "<stdin>", line 1, in <module>
|
|
RuntimeError: unsupported tensor layout: SparseCsr
|
|
>>> b_s.sin()
|
|
tensor(crow_indices=tensor([0, 3, 6]),
|
|
col_indices=tensor([2, 3, 4, 0, 1, 3]),
|
|
values=tensor([ 0.8415, 0.9093, 0.1411, -0.7568, -0.9589, -0.2794]),
|
|
size=(2, 6), nnz=6, layout=torch.sparse_csr)
|
|
|
|
As shown in the example above, we don't support non-zero preserving unary
|
|
operators such as cos. The output of a non-zero preserving unary operation
|
|
will not be able to take advantage of sparse storage formats to the same
|
|
extent as the input and potentially result in a catastrophic increase in memory.
|
|
We instead rely on the user to explicitly convert to a dense Tensor first and
|
|
then run the operation.
|
|
|
|
>>> b_s.to_dense().cos()
|
|
tensor([[ 1.0000, -0.4161],
|
|
[-0.9900, 1.0000]])
|
|
|
|
We are aware that some users want to ignore compressed zeros for operations such
|
|
as `cos` instead of preserving the exact semantics of the operation. For this we
|
|
can point to torch.masked and its MaskedTensor, which is in turn also backed and
|
|
powered by sparse storage formats and kernels.
|
|
|
|
Also note that, for now, the user doesn't have a choice of the output layout. For example,
|
|
adding a sparse Tensor to a regular strided Tensor results in a strided Tensor. Some
|
|
users might prefer for this to stay a sparse layout, because they know the result will
|
|
still be sufficiently sparse.
|
|
|
|
>>> a + b.to_sparse()
|
|
tensor([[0., 3.],
|
|
[3., 0.]])
|
|
|
|
We acknowledge that access to kernels that can efficiently produce different output
|
|
layouts can be very useful. A subsequent operation might significantly benefit from
|
|
receiving a particular layout. We are working on an API to control the result layout
|
|
and recognize it is an important feature to plan a more optimal path of execution for
|
|
any given model.
|
|
|
|
.. _sparse-semi-structured-docs:
|
|
|
|
Sparse Semi-Structured Tensors
|
|
++++++++++++++++++++++++++++++
|
|
|
|
.. warning::
|
|
|
|
Sparse semi-sturctured tensors are currently a prototype feature and subject to change. Please feel free to open an issue to report a bug or if you have feedback to share.
|
|
|
|
Semi-Structured sparsity is a sparse data layout that was first introduced in NVIDIA's Ampere architecture. It is also referred to as **fine-grained structured sparsity** or **2:4 structured sparsity**.
|
|
|
|
This sparse layout stores `n` elements out of every `2n` elements, with `n` being determined by the width of the Tensor's data type (dtype). The most frequently used dtype is float16, where `n=2`, thus the term "2:4 structured sparsity."
|
|
|
|
Semi-structured sparsity is explained in greater detail in `this NVIDIA blog post <https://developer.nvidia.com/blog/exploiting-ampere-structured-sparsity-with-cusparselt>`_.
|
|
|
|
In PyTorch, semi-structured sparsity is implemented via a Tensor subclass.
|
|
By subclassing, we can override ``__torch_dispatch__`` , allowing us to use faster sparse kernels when performing matrix multiplication.
|
|
We can also store the tensor in it's compressed form inside the subclass to reduce memory overhead.
|
|
|
|
In this compressed form, the sparse tensor is stored by retaining only the *specified* elements and some metadata, which encodes the mask.
|
|
|
|
.. note::
|
|
The specified elements and metadata mask of a semi-structured sparse tensor are stored together in a single
|
|
flat compressed tensor. They are appended to each other to form a contiguous chunk of memory.
|
|
|
|
compressed tensor = [ specified elements of original tensor | metadata_mask ]
|
|
|
|
For an original tensor of size `(r, c)` we expect the first `m * k // 2` elements to be the kept elements
|
|
and the rest of the tensor is metadata.
|
|
|
|
In order to make it easier for the user to view the specified elements
|
|
and mask, one can use ``.indices()`` and ``.values()`` to access the mask and specified elements respectively.
|
|
|
|
|
|
- ``.values()`` returns the specified elements in a tensor of size `(r, c//2)` and with the same dtype as the dense matrix.
|
|
|
|
- ``.indices()`` returns the metadata_mask in a tensor of size `(r, c//2 )` and with element type ``torch.int16`` if dtype is torch.float16 or torch.bfloat16, and element type ``torch.int32`` if dtype is torch.int8.
|
|
|
|
|
|
For 2:4 sparse tensors, the metadata overhead is minor - just 2 bits per specified element.
|
|
|
|
.. note::
|
|
It's important to note that ``torch.float32`` is only supported for 1:2 sparsity. Therefore, it does not follow the same formula as above.
|
|
|
|
Here, we break down how to calculate the compression ratio ( size dense / size sparse) of a 2:4 sparse tensor.
|
|
|
|
Let `(r, c) = tensor.shape` and `e = bitwidth(tensor.dtype)`, so `e = 16` for ``torch.float16`` and ``torch.bfloat16`` and `e = 8` for ``torch.int8``.
|
|
|
|
.. math::
|
|
M_{dense} = r \times c \times e \\
|
|
M_{sparse} = M_{specified} + M_{metadata} = r \times \frac{c}{2} \times e + r \times \frac{c}{2} \times 2 = \frac{rce}{2} + rc =rce(\frac{1}{2} +\frac{1}{e})
|
|
|
|
Using these calculations, we can determine the total memory footprint for both the original dense and the new sparse representation.
|
|
|
|
This gives us a simple formula for the compression ratio, which is dependent only on the bitwidth of the tensor datatype.
|
|
|
|
.. math::
|
|
C = \frac{M_{sparse}}{M_{dense}} = \frac{1}{2} + \frac{1}{e}
|
|
|
|
By using this formula, we find that the compression ratio is 56.25% for ``torch.float16`` or ``torch.bfloat16``, and 62.5% for ``torch.int8``.
|
|
|
|
Constructing Sparse Semi-Structured Tensors
|
|
-------------------------------------------
|
|
|
|
You can transform a dense tensor into a sparse semi-structured tensor by simply using the ``torch.to_sparse_semi_structured`` function.
|
|
|
|
Please also note that we only support CUDA tensors since hardware compatibility for semi-structured sparsity is limited to NVIDIA GPUs.
|
|
|
|
|
|
The following datatypes are supported for semi-structured sparsity. Note that each datatype has its own shape constraints and compression factor.
|
|
|
|
.. csv-table::
|
|
:header: "PyTorch dtype", "Shape Constraints", "Compression Factor", "Sparsity Pattern"
|
|
:widths: 15, 45, 10, 10
|
|
:delim: ;
|
|
|
|
``torch.float16``; Tensor must be 2D and (r, c) must both be a positive multiple of 64;9/16;2:4
|
|
``torch.bfloat16``; Tensor must be 2D and (r, c) must both be a positive multiple of 64;9/16;2:4
|
|
``torch.int8``; Tensor must be 2D and (r, c) must both be a positive multiple of 128;10/16;2:4
|
|
|
|
|
|
To construct a semi-structured sparse tensor, start by creating a regular dense tensor that adheres to a 2:4 (or semi-structured) sparse format.
|
|
To do this we tile a small 1x4 strip to create a 16x16 dense float16 tensor.
|
|
Afterwards, we can call ``to_sparse_semi_structured`` function to compress it for accelerated inference.
|
|
|
|
>>> from torch.sparse import to_sparse_semi_structured
|
|
>>> A = torch.Tensor([0, 0, 1, 1]).tile((128, 32)).half().cuda()
|
|
tensor([[0., 0., 1., ..., 0., 1., 1.],
|
|
[0., 0., 1., ..., 0., 1., 1.],
|
|
[0., 0., 1., ..., 0., 1., 1.],
|
|
...,
|
|
[0., 0., 1., ..., 0., 1., 1.],
|
|
[0., 0., 1., ..., 0., 1., 1.],
|
|
[0., 0., 1., ..., 0., 1., 1.]], device='cuda:0', dtype=torch.float16)
|
|
>>> A_sparse = to_sparse_semi_structured(A)
|
|
SparseSemiStructuredTensor(shape=torch.Size([128, 128]), transposed=False, values=tensor([[1., 1., 1., ..., 1., 1., 1.],
|
|
[1., 1., 1., ..., 1., 1., 1.],
|
|
[1., 1., 1., ..., 1., 1., 1.],
|
|
...,
|
|
[1., 1., 1., ..., 1., 1., 1.],
|
|
[1., 1., 1., ..., 1., 1., 1.],
|
|
[1., 1., 1., ..., 1., 1., 1.]], device='cuda:0', dtype=torch.float16), metadata=tensor([[-4370, -4370, -4370, ..., -4370, -4370, -4370],
|
|
[-4370, -4370, -4370, ..., -4370, -4370, -4370],
|
|
[-4370, -4370, -4370, ..., -4370, -4370, -4370],
|
|
...,
|
|
[-4370, -4370, -4370, ..., -4370, -4370, -4370],
|
|
[-4370, -4370, -4370, ..., -4370, -4370, -4370],
|
|
[-4370, -4370, -4370, ..., -4370, -4370, -4370]], device='cuda:0',
|
|
dtype=torch.int16))
|
|
|
|
Sparse Semi-Structured Tensor Operations
|
|
----------------------------------------
|
|
|
|
Currently, the following operations are supported for semi-structured sparse tensors:
|
|
|
|
- torch.addmm(bias, dense, sparse.t())
|
|
- torch.mm(dense, sparse)
|
|
- torch.mm(sparse, dense)
|
|
- aten.linear.default(dense, sparse, bias)
|
|
- aten.t.default(sparse)
|
|
- aten.t.detach(sparse)
|
|
|
|
To use these ops, simply pass the output of ``to_sparse_semi_structured(tensor)`` instead of using ``tensor`` once your tensor has 0s in a semi-structured sparse format, like this:
|
|
|
|
>>> a = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).half().cuda()
|
|
>>> b = torch.rand(64, 64).half().cuda()
|
|
>>> c = torch.mm(a, b)
|
|
>>> a_sparse = to_sparse_semi_structured(a)
|
|
>>> torch.allclose(c, torch.mm(a_sparse, b))
|
|
True
|
|
|
|
Accelerating nn.Linear with semi-structured sparsity
|
|
----------------------------------------------------
|
|
You can accelerate the linear layers in your model if the weights are already semi-structured sparse with just a few lines of code:
|
|
|
|
>>> input = torch.rand(64, 64).half().cuda()
|
|
>>> mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
|
|
>>> linear = nn.Linear(64, 64).half().cuda()
|
|
>>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))
|
|
|
|
|
|
.. _sparse-coo-docs:
|
|
|
|
Sparse COO tensors
|
|
++++++++++++++++++
|
|
|
|
PyTorch implements the so-called Coordinate format, or COO
|
|
format, as one of the storage formats for implementing sparse
|
|
tensors. In COO format, the specified elements are stored as tuples
|
|
of element indices and the corresponding values. In particular,
|
|
|
|
- the indices of specified elements are collected in ``indices``
|
|
tensor of size ``(ndim, nse)`` and with element type
|
|
``torch.int64``,
|
|
|
|
- the corresponding values are collected in ``values`` tensor of
|
|
size ``(nse,)`` and with an arbitrary integer or floating point
|
|
number element type,
|
|
|
|
where ``ndim`` is the dimensionality of the tensor and ``nse`` is the
|
|
number of specified elements.
|
|
|
|
.. note::
|
|
|
|
The memory consumption of a sparse COO tensor is at least ``(ndim *
|
|
8 + <size of element type in bytes>) * nse`` bytes (plus a constant
|
|
overhead from storing other tensor data).
|
|
|
|
The memory consumption of a strided tensor is at least
|
|
``product(<tensor shape>) * <size of element type in bytes>``.
|
|
|
|
For example, the memory consumption of a 10 000 x 10 000 tensor
|
|
with 100 000 non-zero 32-bit floating point numbers is at least
|
|
``(2 * 8 + 4) * 100 000 = 2 000 000`` bytes when using COO tensor
|
|
layout and ``10 000 * 10 000 * 4 = 400 000 000`` bytes when using
|
|
the default strided tensor layout. Notice the 200 fold memory
|
|
saving from using the COO storage format.
|
|
|
|
Construction
|
|
------------
|
|
|
|
A sparse COO tensor can be constructed by providing the two tensors of
|
|
indices and values, as well as the size of the sparse tensor (when it
|
|
cannot be inferred from the indices and values tensors) to a function
|
|
:func:`torch.sparse_coo_tensor`.
|
|
|
|
Suppose we want to define a sparse tensor with the entry 3 at location
|
|
(0, 2), entry 4 at location (1, 0), and entry 5 at location (1, 2).
|
|
Unspecified elements are assumed to have the same value, fill value,
|
|
which is zero by default. We would then write:
|
|
|
|
>>> i = [[0, 1, 1],
|
|
[2, 0, 2]]
|
|
>>> v = [3, 4, 5]
|
|
>>> s = torch.sparse_coo_tensor(i, v, (2, 3))
|
|
>>> s
|
|
tensor(indices=tensor([[0, 1, 1],
|
|
[2, 0, 2]]),
|
|
values=tensor([3, 4, 5]),
|
|
size=(2, 3), nnz=3, layout=torch.sparse_coo)
|
|
>>> s.to_dense()
|
|
tensor([[0, 0, 3],
|
|
[4, 0, 5]])
|
|
|
|
Note that the input ``i`` is NOT a list of index tuples. If you want
|
|
to write your indices this way, you should transpose before passing them to
|
|
the sparse constructor:
|
|
|
|
>>> i = [[0, 2], [1, 0], [1, 2]]
|
|
>>> v = [3, 4, 5 ]
|
|
>>> s = torch.sparse_coo_tensor(list(zip(*i)), v, (2, 3))
|
|
>>> # Or another equivalent formulation to get s
|
|
>>> s = torch.sparse_coo_tensor(torch.tensor(i).t(), v, (2, 3))
|
|
>>> torch.sparse_coo_tensor(i.t(), v, torch.Size([2,3])).to_dense()
|
|
tensor([[0, 0, 3],
|
|
[4, 0, 5]])
|
|
|
|
An empty sparse COO tensor can be constructed by specifying its size
|
|
only:
|
|
|
|
>>> torch.sparse_coo_tensor(size=(2, 3))
|
|
tensor(indices=tensor([], size=(2, 0)),
|
|
values=tensor([], size=(0,)),
|
|
size=(2, 3), nnz=0, layout=torch.sparse_coo)
|
|
|
|
.. _sparse-hybrid-coo-docs:
|
|
|
|
Sparse hybrid COO tensors
|
|
-------------------------
|
|
|
|
PyTorch implements an extension of sparse tensors with scalar values
|
|
to sparse tensors with (contiguous) tensor values. Such tensors are
|
|
called hybrid tensors.
|
|
|
|
PyTorch hybrid COO tensor extends the sparse COO tensor by allowing
|
|
the ``values`` tensor to be a multi-dimensional tensor so that we
|
|
have:
|
|
|
|
- the indices of specified elements are collected in ``indices``
|
|
tensor of size ``(sparse_dims, nse)`` and with element type
|
|
``torch.int64``,
|
|
|
|
- the corresponding (tensor) values are collected in ``values``
|
|
tensor of size ``(nse, dense_dims)`` and with an arbitrary integer
|
|
or floating point number element type.
|
|
|
|
.. note::
|
|
|
|
We use (M + K)-dimensional tensor to denote a N-dimensional sparse
|
|
hybrid tensor, where M and K are the numbers of sparse and dense
|
|
dimensions, respectively, such that M + K == N holds.
|
|
|
|
Suppose we want to create a (2 + 1)-dimensional tensor with the entry
|
|
[3, 4] at location (0, 2), entry [5, 6] at location (1, 0), and entry
|
|
[7, 8] at location (1, 2). We would write
|
|
|
|
>>> i = [[0, 1, 1],
|
|
[2, 0, 2]]
|
|
>>> v = [[3, 4], [5, 6], [7, 8]]
|
|
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
|
|
>>> s
|
|
tensor(indices=tensor([[0, 1, 1],
|
|
[2, 0, 2]]),
|
|
values=tensor([[3, 4],
|
|
[5, 6],
|
|
[7, 8]]),
|
|
size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)
|
|
|
|
>>> s.to_dense()
|
|
tensor([[[0, 0],
|
|
[0, 0],
|
|
[3, 4]],
|
|
[[5, 6],
|
|
[0, 0],
|
|
[7, 8]]])
|
|
|
|
In general, if ``s`` is a sparse COO tensor and ``M =
|
|
s.sparse_dim()``, ``K = s.dense_dim()``, then we have the following
|
|
invariants:
|
|
|
|
- ``M + K == len(s.shape) == s.ndim`` - dimensionality of a tensor
|
|
is the sum of the number of sparse and dense dimensions,
|
|
- ``s.indices().shape == (M, nse)`` - sparse indices are stored
|
|
explicitly,
|
|
- ``s.values().shape == (nse,) + s.shape[M : M + K]`` - the values
|
|
of a hybrid tensor are K-dimensional tensors,
|
|
- ``s.values().layout == torch.strided`` - values are stored as
|
|
strided tensors.
|
|
|
|
.. note::
|
|
|
|
Dense dimensions always follow sparse dimensions, that is, mixing
|
|
of dense and sparse dimensions is not supported.
|
|
|
|
.. note::
|
|
|
|
To be sure that a constructed sparse tensor has consistent indices,
|
|
values, and size, the invariant checks can be enabled per tensor
|
|
creation via ``check_invariants=True`` keyword argument, or
|
|
globally using :class:`torch.sparse.check_sparse_tensor_invariants`
|
|
context manager instance. By default, the sparse tensor invariants
|
|
checks are disabled.
|
|
|
|
.. _sparse-uncoalesced-coo-docs:
|
|
|
|
Uncoalesced sparse COO tensors
|
|
------------------------------
|
|
|
|
PyTorch sparse COO tensor format permits sparse *uncoalesced* tensors,
|
|
where there may be duplicate coordinates in the indices; in this case,
|
|
the interpretation is that the value at that index is the sum of all
|
|
duplicate value entries. For example, one can specify multiple values,
|
|
``3`` and ``4``, for the same index ``1``, that leads to an 1-D
|
|
uncoalesced tensor:
|
|
|
|
>>> i = [[1, 1]]
|
|
>>> v = [3, 4]
|
|
>>> s=torch.sparse_coo_tensor(i, v, (3,))
|
|
>>> s
|
|
tensor(indices=tensor([[1, 1]]),
|
|
values=tensor( [3, 4]),
|
|
size=(3,), nnz=2, layout=torch.sparse_coo)
|
|
|
|
while the coalescing process will accumulate the multi-valued elements
|
|
into a single value using summation:
|
|
|
|
>>> s.coalesce()
|
|
tensor(indices=tensor([[1]]),
|
|
values=tensor([7]),
|
|
size=(3,), nnz=1, layout=torch.sparse_coo)
|
|
|
|
In general, the output of :meth:`torch.Tensor.coalesce` method is a
|
|
sparse tensor with the following properties:
|
|
|
|
- the indices of specified tensor elements are unique,
|
|
- the indices are sorted in lexicographical order,
|
|
- :meth:`torch.Tensor.is_coalesced()` returns ``True``.
|
|
|
|
.. note::
|
|
|
|
For the most part, you shouldn't have to care whether or not a
|
|
sparse tensor is coalesced or not, as most operations will work
|
|
identically given a sparse coalesced or uncoalesced tensor.
|
|
|
|
However, some operations can be implemented more efficiently on
|
|
uncoalesced tensors, and some on coalesced tensors.
|
|
|
|
For instance, addition of sparse COO tensors is implemented by
|
|
simply concatenating the indices and values tensors:
|
|
|
|
>>> a = torch.sparse_coo_tensor([[1, 1]], [5, 6], (2,))
|
|
>>> b = torch.sparse_coo_tensor([[0, 0]], [7, 8], (2,))
|
|
>>> a + b
|
|
tensor(indices=tensor([[0, 0, 1, 1]]),
|
|
values=tensor([7, 8, 5, 6]),
|
|
size=(2,), nnz=4, layout=torch.sparse_coo)
|
|
|
|
If you repeatedly perform an operation that can produce duplicate
|
|
entries (e.g., :func:`torch.Tensor.add`), you should occasionally
|
|
coalesce your sparse tensors to prevent them from growing too large.
|
|
|
|
On the other hand, the lexicographical ordering of indices can be
|
|
advantageous for implementing algorithms that involve many element
|
|
selection operations, such as slicing or matrix products.
|
|
|
|
Working with sparse COO tensors
|
|
-------------------------------
|
|
|
|
Let's consider the following example:
|
|
|
|
>>> i = [[0, 1, 1],
|
|
[2, 0, 2]]
|
|
>>> v = [[3, 4], [5, 6], [7, 8]]
|
|
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
|
|
|
|
As mentioned above, a sparse COO tensor is a :class:`torch.Tensor`
|
|
instance and to distinguish it from the `Tensor` instances that use
|
|
some other layout, on can use :attr:`torch.Tensor.is_sparse` or
|
|
:attr:`torch.Tensor.layout` properties:
|
|
|
|
>>> isinstance(s, torch.Tensor)
|
|
True
|
|
>>> s.is_sparse
|
|
True
|
|
>>> s.layout == torch.sparse_coo
|
|
True
|
|
|
|
The number of sparse and dense dimensions can be acquired using
|
|
methods :meth:`torch.Tensor.sparse_dim` and
|
|
:meth:`torch.Tensor.dense_dim`, respectively. For instance:
|
|
|
|
>>> s.sparse_dim(), s.dense_dim()
|
|
(2, 1)
|
|
|
|
|
|
If ``s`` is a sparse COO tensor then its COO format data can be
|
|
acquired using methods :meth:`torch.Tensor.indices()` and
|
|
:meth:`torch.Tensor.values()`.
|
|
|
|
.. note::
|
|
|
|
Currently, one can acquire the COO format data only when the tensor
|
|
instance is coalesced:
|
|
|
|
>>> s.indices()
|
|
RuntimeError: Cannot get indices on an uncoalesced tensor, please call .coalesce() first
|
|
|
|
For acquiring the COO format data of an uncoalesced tensor, use
|
|
:func:`torch.Tensor._values()` and :func:`torch.Tensor._indices()`:
|
|
|
|
>>> s._indices()
|
|
tensor([[0, 1, 1],
|
|
[2, 0, 2]])
|
|
|
|
.. See https://github.com/pytorch/pytorch/pull/45695 for a new API.
|
|
|
|
.. warning::
|
|
Calling :meth:`torch.Tensor._values()` will return a *detached* tensor.
|
|
To track gradients, :meth:`torch.Tensor.coalesce().values()` must be
|
|
used instead.
|
|
|
|
Constructing a new sparse COO tensor results a tensor that is not
|
|
coalesced:
|
|
|
|
>>> s.is_coalesced()
|
|
False
|
|
|
|
but one can construct a coalesced copy of a sparse COO tensor using
|
|
the :meth:`torch.Tensor.coalesce` method:
|
|
|
|
>>> s2 = s.coalesce()
|
|
>>> s2.indices()
|
|
tensor([[0, 1, 1],
|
|
[2, 0, 2]])
|
|
|
|
When working with uncoalesced sparse COO tensors, one must take into
|
|
an account the additive nature of uncoalesced data: the values of the
|
|
same indices are the terms of a sum that evaluation gives the value of
|
|
the corresponding tensor element. For example, the scalar
|
|
multiplication on a sparse uncoalesced tensor could be implemented by
|
|
multiplying all the uncoalesced values with the scalar because ``c *
|
|
(a + b) == c * a + c * b`` holds. However, any nonlinear operation,
|
|
say, a square root, cannot be implemented by applying the operation to
|
|
uncoalesced data because ``sqrt(a + b) == sqrt(a) + sqrt(b)`` does not
|
|
hold in general.
|
|
|
|
Slicing (with positive step) of a sparse COO tensor is supported only
|
|
for dense dimensions. Indexing is supported for both sparse and dense
|
|
dimensions:
|
|
|
|
>>> s[1]
|
|
tensor(indices=tensor([[0, 2]]),
|
|
values=tensor([[5, 6],
|
|
[7, 8]]),
|
|
size=(3, 2), nnz=2, layout=torch.sparse_coo)
|
|
>>> s[1, 0, 1]
|
|
tensor(6)
|
|
>>> s[1, 0, 1:]
|
|
tensor([6])
|
|
|
|
|
|
In PyTorch, the fill value of a sparse tensor cannot be specified
|
|
explicitly and is assumed to be zero in general. However, there exists
|
|
operations that may interpret the fill value differently. For
|
|
instance, :func:`torch.sparse.softmax` computes the softmax with the
|
|
assumption that the fill value is negative infinity.
|
|
|
|
.. See https://github.com/Quansight-Labs/rfcs/tree/pearu/rfc-fill-value/RFC-0004-sparse-fill-value for a new API
|
|
|
|
.. _sparse-compressed-docs:
|
|
|
|
Sparse Compressed Tensors
|
|
+++++++++++++++++++++++++
|
|
|
|
Sparse Compressed Tensors represents a class of sparse tensors that
|
|
have a common feature of compressing the indices of a certain dimension
|
|
using an encoding that enables certain optimizations on linear algebra
|
|
kernels of sparse compressed tensors. This encoding is based on the
|
|
`Compressed Sparse Row (CSR)`__ format that PyTorch sparse compressed
|
|
tensors extend with the support of sparse tensor batches, allowing
|
|
multi-dimensional tensor values, and storing sparse tensor values in
|
|
dense blocks.
|
|
|
|
__ https://en.wikipedia.org/wiki/Sparse_matrix#Compressed_sparse_row_(CSR,_CRS_or_Yale_format)
|
|
|
|
.. note::
|
|
|
|
We use (B + M + K)-dimensional tensor to denote a N-dimensional
|
|
sparse compressed hybrid tensor, where B, M, and K are the numbers
|
|
of batch, sparse, and dense dimensions, respectively, such that
|
|
``B + M + K == N`` holds. The number of sparse dimensions for
|
|
sparse compressed tensors is always two, ``M == 2``.
|
|
|
|
.. note::
|
|
|
|
We say that an indices tensor ``compressed_indices`` uses CSR
|
|
encoding if the following invariants are satisfied:
|
|
|
|
- ``compressed_indices`` is a contiguous strided 32 or 64 bit
|
|
integer tensor
|
|
- ``compressed_indices`` shape is ``(*batchsize,
|
|
compressed_dim_size + 1)`` where ``compressed_dim_size`` is the
|
|
number of compressed dimensions (e.g. rows or columns)
|
|
- ``compressed_indices[..., 0] == 0`` where ``...`` denotes batch
|
|
indices
|
|
- ``compressed_indices[..., compressed_dim_size] == nse`` where
|
|
``nse`` is the number of specified elements
|
|
- ``0 <= compressed_indices[..., i] - compressed_indices[..., i -
|
|
1] <= plain_dim_size`` for ``i=1, ..., compressed_dim_size``,
|
|
where ``plain_dim_size`` is the number of plain dimensions
|
|
(orthogonal to compressed dimensions, e.g. columns or rows).
|
|
|
|
To be sure that a constructed sparse tensor has consistent indices,
|
|
values, and size, the invariant checks can be enabled per tensor
|
|
creation via ``check_invariants=True`` keyword argument, or
|
|
globally using :class:`torch.sparse.check_sparse_tensor_invariants`
|
|
context manager instance. By default, the sparse tensor invariants
|
|
checks are disabled.
|
|
|
|
.. note::
|
|
|
|
The generalization of sparse compressed layouts to N-dimensional
|
|
tensors can lead to some confusion regarding the count of specified
|
|
elements. When a sparse compressed tensor contains batch dimensions
|
|
the number of specified elements will correspond to the number of such
|
|
elements per-batch. When a sparse compressed tensor has dense dimensions
|
|
the element considered is now the K-dimensional array. Also for block
|
|
sparse compressed layouts the 2-D block is considered as the element
|
|
being specified. Take as an example a 3-dimensional block sparse
|
|
tensor, with one batch dimension of length ``b``, and a block
|
|
shape of ``p, q``. If this tensor has ``n`` specified elements, then
|
|
in fact we have ``n`` blocks specified per batch. This tensor would
|
|
have ``values`` with shape ``(b, n, p, q)``. This interpretation of the
|
|
number of specified elements comes from all sparse compressed layouts
|
|
being derived from the compression of a 2-dimensional matrix. Batch
|
|
dimensions are treated as stacking of sparse matrices, dense dimensions
|
|
change the meaning of the element from a simple scalar value to an
|
|
array with its own dimensions.
|
|
|
|
.. _sparse-csr-docs:
|
|
|
|
Sparse CSR Tensor
|
|
-----------------
|
|
|
|
The primary advantage of the CSR format over the COO format is better
|
|
use of storage and much faster computation operations such as sparse
|
|
matrix-vector multiplication using MKL and MAGMA backends.
|
|
|
|
In the simplest case, a (0 + 2 + 0)-dimensional sparse CSR tensor
|
|
consists of three 1-D tensors: ``crow_indices``, ``col_indices`` and
|
|
``values``:
|
|
|
|
- The ``crow_indices`` tensor consists of compressed row
|
|
indices. This is a 1-D tensor of size ``nrows + 1`` (the number of
|
|
rows plus 1). The last element of ``crow_indices`` is the number
|
|
of specified elements, ``nse``. This tensor encodes the index in
|
|
``values`` and ``col_indices`` depending on where the given row
|
|
starts. Each successive number in the tensor subtracted by the
|
|
number before it denotes the number of elements in a given row.
|
|
|
|
- The ``col_indices`` tensor contains the column indices of each
|
|
element. This is a 1-D tensor of size ``nse``.
|
|
|
|
- The ``values`` tensor contains the values of the CSR tensor
|
|
elements. This is a 1-D tensor of size ``nse``.
|
|
|
|
.. note::
|
|
|
|
The index tensors ``crow_indices`` and ``col_indices`` should have
|
|
element type either ``torch.int64`` (default) or
|
|
``torch.int32``. If you want to use MKL-enabled matrix operations,
|
|
use ``torch.int32``. This is as a result of the default linking of
|
|
pytorch being with MKL LP64, which uses 32 bit integer indexing.
|
|
|
|
In the general case, the (B + 2 + K)-dimensional sparse CSR tensor
|
|
consists of two (B + 1)-dimensional index tensors ``crow_indices`` and
|
|
``col_indices``, and of (1 + K)-dimensional ``values`` tensor such
|
|
that
|
|
|
|
- ``crow_indices.shape == (*batchsize, nrows + 1)``
|
|
|
|
- ``col_indices.shape == (*batchsize, nse)``
|
|
|
|
- ``values.shape == (nse, *densesize)``
|
|
|
|
while the shape of the sparse CSR tensor is ``(*batchsize, nrows,
|
|
ncols, *densesize)`` where ``len(batchsize) == B`` and
|
|
``len(densesize) == K``.
|
|
|
|
.. note::
|
|
|
|
The batches of sparse CSR tensors are dependent: the number of
|
|
specified elements in all batches must be the same. This somewhat
|
|
artificial constraint allows efficient storage of the indices of
|
|
different CSR batches.
|
|
|
|
.. note::
|
|
|
|
The number of sparse and dense dimensions can be acquired using
|
|
:meth:`torch.Tensor.sparse_dim` and :meth:`torch.Tensor.dense_dim`
|
|
methods. The batch dimensions can be computed from the tensor
|
|
shape: ``batchsize = tensor.shape[:-tensor.sparse_dim() -
|
|
tensor.dense_dim()]``.
|
|
|
|
.. note::
|
|
|
|
The memory consumption of a sparse CSR tensor is at least
|
|
``(nrows * 8 + (8 + <size of element type in bytes> *
|
|
prod(densesize)) * nse) * prod(batchsize)`` bytes (plus a constant
|
|
overhead from storing other tensor data).
|
|
|
|
With the same example data of :ref:`the note in sparse COO format
|
|
introduction<sparse-coo-docs>`, the memory consumption of a 10 000
|
|
x 10 000 tensor with 100 000 non-zero 32-bit floating point numbers
|
|
is at least ``(10000 * 8 + (8 + 4 * 1) * 100 000) * 1 = 1 280 000``
|
|
bytes when using CSR tensor layout. Notice the 1.6 and 310 fold
|
|
savings from using CSR storage format compared to using the COO and
|
|
strided formats, respectively.
|
|
|
|
Construction of CSR tensors
|
|
'''''''''''''''''''''''''''
|
|
|
|
Sparse CSR tensors can be directly constructed by using the
|
|
:func:`torch.sparse_csr_tensor` function. The user must supply the row
|
|
and column indices and values tensors separately where the row indices
|
|
must be specified using the CSR compression encoding. The ``size``
|
|
argument is optional and will be deduced from the ``crow_indices`` and
|
|
``col_indices`` if it is not present.
|
|
|
|
>>> crow_indices = torch.tensor([0, 2, 4])
|
|
>>> col_indices = torch.tensor([0, 1, 0, 1])
|
|
>>> values = torch.tensor([1, 2, 3, 4])
|
|
>>> csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
|
|
>>> csr
|
|
tensor(crow_indices=tensor([0, 2, 4]),
|
|
col_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
|
|
dtype=torch.float64)
|
|
>>> csr.to_dense()
|
|
tensor([[1., 2.],
|
|
[3., 4.]], dtype=torch.float64)
|
|
|
|
.. note::
|
|
|
|
The values of sparse dimensions in deduced ``size`` is computed
|
|
from the size of ``crow_indices`` and the maximal index value in
|
|
``col_indices``. If the number of columns needs to be larger than
|
|
in the deduced ``size`` then the ``size`` argument must be
|
|
specified explicitly.
|
|
|
|
The simplest way of constructing a 2-D sparse CSR tensor from a
|
|
strided or sparse COO tensor is to use
|
|
:meth:`torch.Tensor.to_sparse_csr` method. Any zeros in the (strided)
|
|
tensor will be interpreted as missing values in the sparse tensor:
|
|
|
|
>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
|
|
>>> sp = a.to_sparse_csr()
|
|
>>> sp
|
|
tensor(crow_indices=tensor([0, 1, 3, 3]),
|
|
col_indices=tensor([2, 0, 1]),
|
|
values=tensor([1., 1., 2.]), size=(3, 4), nnz=3, dtype=torch.float64)
|
|
|
|
CSR Tensor Operations
|
|
'''''''''''''''''''''
|
|
|
|
The sparse matrix-vector multiplication can be performed with the
|
|
:meth:`tensor.matmul` method. This is currently the only math operation
|
|
supported on CSR tensors.
|
|
|
|
>>> vec = torch.randn(4, 1, dtype=torch.float64)
|
|
>>> sp.matmul(vec)
|
|
tensor([[0.9078],
|
|
[1.3180],
|
|
[0.0000]], dtype=torch.float64)
|
|
|
|
.. _sparse-csc-docs:
|
|
|
|
Sparse CSC Tensor
|
|
-----------------
|
|
|
|
The sparse CSC (Compressed Sparse Column) tensor format implements the
|
|
CSC format for storage of 2 dimensional tensors with an extension to
|
|
supporting batches of sparse CSC tensors and values being
|
|
multi-dimensional tensors.
|
|
|
|
.. note::
|
|
|
|
Sparse CSC tensor is essentially a transpose of the sparse CSR
|
|
tensor when the transposition is about swapping the sparse
|
|
dimensions.
|
|
|
|
Similarly to :ref:`sparse CSR tensors <sparse-csr-docs>`, a sparse CSC
|
|
tensor consists of three tensors: ``ccol_indices``, ``row_indices``
|
|
and ``values``:
|
|
|
|
- The ``ccol_indices`` tensor consists of compressed column
|
|
indices. This is a (B + 1)-D tensor of shape ``(*batchsize, ncols + 1)``.
|
|
The last element is the number of specified
|
|
elements, ``nse``. This tensor encodes the index in ``values`` and
|
|
``row_indices`` depending on where the given column starts. Each
|
|
successive number in the tensor subtracted by the number before it
|
|
denotes the number of elements in a given column.
|
|
|
|
- The ``row_indices`` tensor contains the row indices of each
|
|
element. This is a (B + 1)-D tensor of shape ``(*batchsize, nse)``.
|
|
|
|
- The ``values`` tensor contains the values of the CSC tensor
|
|
elements. This is a (1 + K)-D tensor of shape ``(nse, *densesize)``.
|
|
|
|
Construction of CSC tensors
|
|
'''''''''''''''''''''''''''
|
|
|
|
Sparse CSC tensors can be directly constructed by using the
|
|
:func:`torch.sparse_csc_tensor` function. The user must supply the row
|
|
and column indices and values tensors separately where the column indices
|
|
must be specified using the CSR compression encoding. The ``size``
|
|
argument is optional and will be deduced from the ``row_indices`` and
|
|
``ccol_indices`` tensors if it is not present.
|
|
|
|
>>> ccol_indices = torch.tensor([0, 2, 4])
|
|
>>> row_indices = torch.tensor([0, 1, 0, 1])
|
|
>>> values = torch.tensor([1, 2, 3, 4])
|
|
>>> csc = torch.sparse_csc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
|
|
>>> csc
|
|
tensor(ccol_indices=tensor([0, 2, 4]),
|
|
row_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
|
|
dtype=torch.float64, layout=torch.sparse_csc)
|
|
>>> csc.to_dense()
|
|
tensor([[1., 3.],
|
|
[2., 4.]], dtype=torch.float64)
|
|
|
|
.. note::
|
|
|
|
The sparse CSC tensor constructor function has the compressed
|
|
column indices argument before the row indices argument.
|
|
|
|
The (0 + 2 + 0)-dimensional sparse CSC tensors can be constructed from
|
|
any two-dimensional tensor using :meth:`torch.Tensor.to_sparse_csc`
|
|
method. Any zeros in the (strided) tensor will be interpreted as
|
|
missing values in the sparse tensor:
|
|
|
|
>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
|
|
>>> sp = a.to_sparse_csc()
|
|
>>> sp
|
|
tensor(ccol_indices=tensor([0, 1, 2, 3, 3]),
|
|
row_indices=tensor([1, 1, 0]),
|
|
values=tensor([1., 2., 1.]), size=(3, 4), nnz=3, dtype=torch.float64,
|
|
layout=torch.sparse_csc)
|
|
|
|
.. _sparse-bsr-docs:
|
|
|
|
Sparse BSR Tensor
|
|
-----------------
|
|
|
|
The sparse BSR (Block compressed Sparse Row) tensor format implements the
|
|
BSR format for storage of two-dimensional tensors with an extension to
|
|
supporting batches of sparse BSR tensors and values being blocks of
|
|
multi-dimensional tensors.
|
|
|
|
A sparse BSR tensor consists of three tensors: ``crow_indices``,
|
|
``col_indices`` and ``values``:
|
|
|
|
- The ``crow_indices`` tensor consists of compressed row
|
|
indices. This is a (B + 1)-D tensor of shape ``(*batchsize,
|
|
nrowblocks + 1)``. The last element is the number of specified blocks,
|
|
``nse``. This tensor encodes the index in ``values`` and
|
|
``col_indices`` depending on where the given column block
|
|
starts. Each successive number in the tensor subtracted by the
|
|
number before it denotes the number of blocks in a given row.
|
|
|
|
- The ``col_indices`` tensor contains the column block indices of each
|
|
element. This is a (B + 1)-D tensor of shape ``(*batchsize,
|
|
nse)``.
|
|
|
|
- The ``values`` tensor contains the values of the sparse BSR tensor
|
|
elements collected into two-dimensional blocks. This is a (1 + 2 +
|
|
K)-D tensor of shape ``(nse, nrowblocks, ncolblocks,
|
|
*densesize)``.
|
|
|
|
Construction of BSR tensors
|
|
'''''''''''''''''''''''''''
|
|
|
|
Sparse BSR tensors can be directly constructed by using the
|
|
:func:`torch.sparse_bsr_tensor` function. The user must supply the row
|
|
and column block indices and values tensors separately where the row block indices
|
|
must be specified using the CSR compression encoding.
|
|
The ``size`` argument is optional and will be deduced from the ``crow_indices`` and
|
|
``col_indices`` tensors if it is not present.
|
|
|
|
>>> crow_indices = torch.tensor([0, 2, 4])
|
|
>>> col_indices = torch.tensor([0, 1, 0, 1])
|
|
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
|
|
... [[3, 4, 5], [9, 10, 11]],
|
|
... [[12, 13, 14], [18, 19, 20]],
|
|
... [[15, 16, 17], [21, 22, 23]]])
|
|
>>> bsr = torch.sparse_bsr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
|
|
>>> bsr
|
|
tensor(crow_indices=tensor([0, 2, 4]),
|
|
col_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([[[ 0., 1., 2.],
|
|
[ 6., 7., 8.]],
|
|
[[ 3., 4., 5.],
|
|
[ 9., 10., 11.]],
|
|
[[12., 13., 14.],
|
|
[18., 19., 20.]],
|
|
[[15., 16., 17.],
|
|
[21., 22., 23.]]]),
|
|
size=(4, 6), nnz=4, dtype=torch.float64, layout=torch.sparse_bsr)
|
|
>>> bsr.to_dense()
|
|
tensor([[ 0., 1., 2., 3., 4., 5.],
|
|
[ 6., 7., 8., 9., 10., 11.],
|
|
[12., 13., 14., 15., 16., 17.],
|
|
[18., 19., 20., 21., 22., 23.]], dtype=torch.float64)
|
|
|
|
The (0 + 2 + 0)-dimensional sparse BSR tensors can be constructed from
|
|
any two-dimensional tensor using :meth:`torch.Tensor.to_sparse_bsr`
|
|
method that also requires the specification of the values block size:
|
|
|
|
>>> dense = torch.tensor([[0, 1, 2, 3, 4, 5],
|
|
... [6, 7, 8, 9, 10, 11],
|
|
... [12, 13, 14, 15, 16, 17],
|
|
... [18, 19, 20, 21, 22, 23]])
|
|
>>> bsr = dense.to_sparse_bsr(blocksize=(2, 3))
|
|
>>> bsr
|
|
tensor(crow_indices=tensor([0, 2, 4]),
|
|
col_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([[[ 0, 1, 2],
|
|
[ 6, 7, 8]],
|
|
[[ 3, 4, 5],
|
|
[ 9, 10, 11]],
|
|
[[12, 13, 14],
|
|
[18, 19, 20]],
|
|
[[15, 16, 17],
|
|
[21, 22, 23]]]), size=(4, 6), nnz=4,
|
|
layout=torch.sparse_bsr)
|
|
|
|
.. _sparse-bsc-docs:
|
|
|
|
Sparse BSC Tensor
|
|
-----------------
|
|
|
|
The sparse BSC (Block compressed Sparse Column) tensor format implements the
|
|
BSC format for storage of two-dimensional tensors with an extension to
|
|
supporting batches of sparse BSC tensors and values being blocks of
|
|
multi-dimensional tensors.
|
|
|
|
A sparse BSC tensor consists of three tensors: ``ccol_indices``,
|
|
``row_indices`` and ``values``:
|
|
|
|
- The ``ccol_indices`` tensor consists of compressed column
|
|
indices. This is a (B + 1)-D tensor of shape ``(*batchsize,
|
|
ncolblocks + 1)``. The last element is the number of specified blocks,
|
|
``nse``. This tensor encodes the index in ``values`` and
|
|
``row_indices`` depending on where the given row block
|
|
starts. Each successive number in the tensor subtracted by the
|
|
number before it denotes the number of blocks in a given column.
|
|
|
|
- The ``row_indices`` tensor contains the row block indices of each
|
|
element. This is a (B + 1)-D tensor of shape ``(*batchsize,
|
|
nse)``.
|
|
|
|
- The ``values`` tensor contains the values of the sparse BSC tensor
|
|
elements collected into two-dimensional blocks. This is a (1 + 2 +
|
|
K)-D tensor of shape ``(nse, nrowblocks, ncolblocks,
|
|
*densesize)``.
|
|
|
|
Construction of BSC tensors
|
|
'''''''''''''''''''''''''''
|
|
|
|
Sparse BSC tensors can be directly constructed by using the
|
|
:func:`torch.sparse_bsc_tensor` function. The user must supply the row
|
|
and column block indices and values tensors separately where the column block indices
|
|
must be specified using the CSR compression encoding.
|
|
The ``size`` argument is optional and will be deduced from the ``ccol_indices`` and
|
|
``row_indices`` tensors if it is not present.
|
|
|
|
>>> ccol_indices = torch.tensor([0, 2, 4])
|
|
>>> row_indices = torch.tensor([0, 1, 0, 1])
|
|
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
|
|
... [[3, 4, 5], [9, 10, 11]],
|
|
... [[12, 13, 14], [18, 19, 20]],
|
|
... [[15, 16, 17], [21, 22, 23]]])
|
|
>>> bsc = torch.sparse_bsc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
|
|
>>> bsc
|
|
tensor(ccol_indices=tensor([0, 2, 4]),
|
|
row_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([[[ 0., 1., 2.],
|
|
[ 6., 7., 8.]],
|
|
[[ 3., 4., 5.],
|
|
[ 9., 10., 11.]],
|
|
[[12., 13., 14.],
|
|
[18., 19., 20.]],
|
|
[[15., 16., 17.],
|
|
[21., 22., 23.]]]), size=(4, 6), nnz=4,
|
|
dtype=torch.float64, layout=torch.sparse_bsc)
|
|
|
|
Tools for working with sparse compressed tensors
|
|
------------------------------------------------
|
|
|
|
All sparse compressed tensors --- CSR, CSC, BSR, and BSC tensors ---
|
|
are conceptionally very similar in that their indices data is split
|
|
into two parts: so-called compressed indices that use the CSR
|
|
encoding, and so-called plain indices that are orthogonal to the
|
|
compressed indices. This allows various tools on these tensors to
|
|
share the same implementations that are parameterized by tensor
|
|
layout.
|
|
|
|
Construction of sparse compressed tensors
|
|
'''''''''''''''''''''''''''''''''''''''''
|
|
|
|
Sparse CSR, CSC, BSR, and CSC tensors can be constructed by using
|
|
:func:`torch.sparse_compressed_tensor` function that have the same
|
|
interface as the above discussed constructor functions
|
|
:func:`torch.sparse_csr_tensor`, :func:`torch.sparse_csc_tensor`,
|
|
:func:`torch.sparse_bsr_tensor`, and :func:`torch.sparse_bsc_tensor`,
|
|
respectively, but with an extra required ``layout`` argument. The
|
|
following example illustrates a method of constructing CSR and CSC
|
|
tensors using the same input data by specifying the corresponding
|
|
layout parameter to the :func:`torch.sparse_compressed_tensor`
|
|
function:
|
|
|
|
>>> compressed_indices = torch.tensor([0, 2, 4])
|
|
>>> plain_indices = torch.tensor([0, 1, 0, 1])
|
|
>>> values = torch.tensor([1, 2, 3, 4])
|
|
>>> csr = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csr)
|
|
>>> csr
|
|
tensor(crow_indices=tensor([0, 2, 4]),
|
|
col_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
|
|
layout=torch.sparse_csr)
|
|
>>> csc = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csc)
|
|
>>> csc
|
|
tensor(ccol_indices=tensor([0, 2, 4]),
|
|
row_indices=tensor([0, 1, 0, 1]),
|
|
values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
|
|
layout=torch.sparse_csc)
|
|
>>> (csr.transpose(0, 1).to_dense() == csc.to_dense()).all()
|
|
tensor(True)
|
|
|
|
.. _sparse-ops-docs:
|
|
|
|
Supported operations
|
|
+++++++++++++++++++++++++++++++++++
|
|
|
|
Linear Algebra operations
|
|
-------------------------
|
|
|
|
The following table summarizes supported Linear Algebra operations on
|
|
sparse matrices where the operands layouts may vary. Here
|
|
``T[layout]`` denotes a tensor with a given layout. Similarly,
|
|
``M[layout]`` denotes a matrix (2-D PyTorch tensor), and ``V[layout]``
|
|
denotes a vector (1-D PyTorch tensor). In addition, ``f`` denotes a
|
|
scalar (float or 0-D PyTorch tensor), ``*`` is element-wise
|
|
multiplication, and ``@`` is matrix multiplication.
|
|
|
|
.. csv-table::
|
|
:header: "PyTorch operation", "Sparse grad?", "Layout signature"
|
|
:widths: 20, 5, 60
|
|
:delim: ;
|
|
|
|
:func:`torch.mv`;no; ``M[sparse_coo] @ V[strided] -> V[strided]``
|
|
:func:`torch.mv`;no; ``M[sparse_csr] @ V[strided] -> V[strided]``
|
|
:func:`torch.matmul`; no; ``M[sparse_coo] @ M[strided] -> M[strided]``
|
|
:func:`torch.matmul`; no; ``M[sparse_csr] @ M[strided] -> M[strided]``
|
|
:func:`torch.matmul`; no; ``M[SparseSemiStructured] @ M[strided] -> M[strided]``
|
|
:func:`torch.matmul`; no; ``M[strided] @ M[SparseSemiStructured] -> M[strided]``
|
|
:func:`torch.mm`; no; ``M[sparse_coo] @ M[strided] -> M[strided]``
|
|
:func:`torch.mm`; no; ``M[SparseSemiStructured] @ M[strided] -> M[strided]``
|
|
:func:`torch.mm`; no; ``M[strided] @ M[SparseSemiStructured] -> M[strided]``
|
|
:func:`torch.sparse.mm`; yes; ``M[sparse_coo] @ M[strided] -> M[strided]``
|
|
:func:`torch.smm`; no; ``M[sparse_coo] @ M[strided] -> M[sparse_coo]``
|
|
:func:`torch.hspmm`; no; ``M[sparse_coo] @ M[strided] -> M[hybrid sparse_coo]``
|
|
:func:`torch.bmm`; no; ``T[sparse_coo] @ T[strided] -> T[strided]``
|
|
:func:`torch.addmm`; no; ``f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]``
|
|
:func:`torch.addmm`; no; ``f * M[strided] + f * (M[SparseSemiStructured] @ M[strided]) -> M[strided]``
|
|
:func:`torch.addmm`; no; ``f * M[strided] + f * (M[strided] @ M[SparseSemiStructured]) -> M[strided]``
|
|
:func:`torch.sparse.addmm`; yes; ``f * M[strided] + f * (M[sparse_coo] @ M[strided]) -> M[strided]``
|
|
:func:`torch.sspaddmm`; no; ``f * M[sparse_coo] + f * (M[sparse_coo] @ M[strided]) -> M[sparse_coo]``
|
|
:func:`torch.lobpcg`; no; ``GENEIG(M[sparse_coo]) -> M[strided], M[strided]``
|
|
:func:`torch.pca_lowrank`; yes; ``PCA(M[sparse_coo]) -> M[strided], M[strided], M[strided]``
|
|
:func:`torch.svd_lowrank`; yes; ``SVD(M[sparse_coo]) -> M[strided], M[strided], M[strided]``
|
|
|
|
where "Sparse grad?" column indicates if the PyTorch operation supports
|
|
backward with respect to sparse matrix argument. All PyTorch operations,
|
|
except :func:`torch.smm`, support backward with respect to strided
|
|
matrix arguments.
|
|
|
|
.. note::
|
|
|
|
Currently, PyTorch does not support matrix multiplication with the
|
|
layout signature ``M[strided] @ M[sparse_coo]``. However,
|
|
applications can still compute this using the matrix relation ``D @
|
|
S == (S.t() @ D.t()).t()``.
|
|
|
|
Tensor methods and sparse
|
|
-------------------------
|
|
|
|
The following Tensor methods are related to sparse tensors:
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
Tensor.is_sparse
|
|
Tensor.is_sparse_csr
|
|
Tensor.dense_dim
|
|
Tensor.sparse_dim
|
|
Tensor.sparse_mask
|
|
Tensor.to_sparse
|
|
Tensor.to_sparse_coo
|
|
Tensor.to_sparse_csr
|
|
Tensor.to_sparse_csc
|
|
Tensor.to_sparse_bsr
|
|
Tensor.to_sparse_bsc
|
|
Tensor.to_dense
|
|
Tensor.values
|
|
|
|
The following Tensor methods are specific to sparse COO tensors:
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
Tensor.coalesce
|
|
Tensor.sparse_resize_
|
|
Tensor.sparse_resize_and_clear_
|
|
Tensor.is_coalesced
|
|
Tensor.indices
|
|
|
|
The following methods are specific to :ref:`sparse CSR tensors <sparse-csr-docs>` and :ref:`sparse BSR tensors <sparse-bsr-docs>`:
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
Tensor.crow_indices
|
|
Tensor.col_indices
|
|
|
|
The following methods are specific to :ref:`sparse CSC tensors <sparse-csc-docs>` and :ref:`sparse BSC tensors <sparse-bsc-docs>`:
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
Tensor.row_indices
|
|
Tensor.ccol_indices
|
|
|
|
The following Tensor methods support sparse COO tensors:
|
|
|
|
:meth:`~torch.Tensor.add`
|
|
:meth:`~torch.Tensor.add_`
|
|
:meth:`~torch.Tensor.addmm`
|
|
:meth:`~torch.Tensor.addmm_`
|
|
:meth:`~torch.Tensor.any`
|
|
:meth:`~torch.Tensor.asin`
|
|
:meth:`~torch.Tensor.asin_`
|
|
:meth:`~torch.Tensor.arcsin`
|
|
:meth:`~torch.Tensor.arcsin_`
|
|
:meth:`~torch.Tensor.bmm`
|
|
:meth:`~torch.Tensor.clone`
|
|
:meth:`~torch.Tensor.deg2rad`
|
|
:meth:`~torch.Tensor.deg2rad_`
|
|
:meth:`~torch.Tensor.detach`
|
|
:meth:`~torch.Tensor.detach_`
|
|
:meth:`~torch.Tensor.dim`
|
|
:meth:`~torch.Tensor.div`
|
|
:meth:`~torch.Tensor.div_`
|
|
:meth:`~torch.Tensor.floor_divide`
|
|
:meth:`~torch.Tensor.floor_divide_`
|
|
:meth:`~torch.Tensor.get_device`
|
|
:meth:`~torch.Tensor.index_select`
|
|
:meth:`~torch.Tensor.isnan`
|
|
:meth:`~torch.Tensor.log1p`
|
|
:meth:`~torch.Tensor.log1p_`
|
|
:meth:`~torch.Tensor.mm`
|
|
:meth:`~torch.Tensor.mul`
|
|
:meth:`~torch.Tensor.mul_`
|
|
:meth:`~torch.Tensor.mv`
|
|
:meth:`~torch.Tensor.narrow_copy`
|
|
:meth:`~torch.Tensor.neg`
|
|
:meth:`~torch.Tensor.neg_`
|
|
:meth:`~torch.Tensor.negative`
|
|
:meth:`~torch.Tensor.negative_`
|
|
:meth:`~torch.Tensor.numel`
|
|
:meth:`~torch.Tensor.rad2deg`
|
|
:meth:`~torch.Tensor.rad2deg_`
|
|
:meth:`~torch.Tensor.resize_as_`
|
|
:meth:`~torch.Tensor.size`
|
|
:meth:`~torch.Tensor.pow`
|
|
:meth:`~torch.Tensor.sqrt`
|
|
:meth:`~torch.Tensor.square`
|
|
:meth:`~torch.Tensor.smm`
|
|
:meth:`~torch.Tensor.sspaddmm`
|
|
:meth:`~torch.Tensor.sub`
|
|
:meth:`~torch.Tensor.sub_`
|
|
:meth:`~torch.Tensor.t`
|
|
:meth:`~torch.Tensor.t_`
|
|
:meth:`~torch.Tensor.transpose`
|
|
:meth:`~torch.Tensor.transpose_`
|
|
:meth:`~torch.Tensor.zero_`
|
|
|
|
Torch functions specific to sparse Tensors
|
|
------------------------------------------
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
sparse_coo_tensor
|
|
sparse_csr_tensor
|
|
sparse_csc_tensor
|
|
sparse_bsr_tensor
|
|
sparse_bsc_tensor
|
|
sparse_compressed_tensor
|
|
sparse.sum
|
|
sparse.addmm
|
|
sparse.sampled_addmm
|
|
sparse.mm
|
|
sspaddmm
|
|
hspmm
|
|
smm
|
|
sparse.softmax
|
|
sparse.log_softmax
|
|
sparse.spdiags
|
|
|
|
Other functions
|
|
---------------
|
|
|
|
The following :mod:`torch` functions support sparse tensors:
|
|
|
|
:func:`~torch.cat`
|
|
:func:`~torch.dstack`
|
|
:func:`~torch.empty`
|
|
:func:`~torch.empty_like`
|
|
:func:`~torch.hstack`
|
|
:func:`~torch.index_select`
|
|
:func:`~torch.is_complex`
|
|
:func:`~torch.is_floating_point`
|
|
:func:`~torch.is_nonzero`
|
|
:func:`~torch.is_same_size`
|
|
:func:`~torch.is_signed`
|
|
:func:`~torch.is_tensor`
|
|
:func:`~torch.lobpcg`
|
|
:func:`~torch.mm`
|
|
:func:`~torch.native_norm`
|
|
:func:`~torch.pca_lowrank`
|
|
:func:`~torch.select`
|
|
:func:`~torch.stack`
|
|
:func:`~torch.svd_lowrank`
|
|
:func:`~torch.unsqueeze`
|
|
:func:`~torch.vstack`
|
|
:func:`~torch.zeros`
|
|
:func:`~torch.zeros_like`
|
|
|
|
To manage checking sparse tensor invariants, see:
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
sparse.check_sparse_tensor_invariants
|
|
|
|
To use sparse tensors with :func:`~torch.autograd.gradcheck` function,
|
|
see:
|
|
|
|
.. autosummary::
|
|
:toctree: generated
|
|
:nosignatures:
|
|
|
|
sparse.as_sparse_gradcheck
|
|
|
|
Unary functions
|
|
---------------
|
|
|
|
We aim to support all zero-preserving unary functions.
|
|
|
|
If you find that we are missing a zero-preserving unary function
|
|
that you need, please feel encouraged to open an issue for a feature request.
|
|
As always please kindly try the search function first before opening an issue.
|
|
|
|
The following operators currently support sparse COO/CSR/CSC/BSR/CSR tensor inputs.
|
|
|
|
:func:`~torch.abs`
|
|
:func:`~torch.asin`
|
|
:func:`~torch.asinh`
|
|
:func:`~torch.atan`
|
|
:func:`~torch.atanh`
|
|
:func:`~torch.ceil`
|
|
:func:`~torch.conj_physical`
|
|
:func:`~torch.floor`
|
|
:func:`~torch.log1p`
|
|
:func:`~torch.neg`
|
|
:func:`~torch.round`
|
|
:func:`~torch.sin`
|
|
:func:`~torch.sinh`
|
|
:func:`~torch.sign`
|
|
:func:`~torch.sgn`
|
|
:func:`~torch.signbit`
|
|
:func:`~torch.tan`
|
|
:func:`~torch.tanh`
|
|
:func:`~torch.trunc`
|
|
:func:`~torch.expm1`
|
|
:func:`~torch.sqrt`
|
|
:func:`~torch.angle`
|
|
:func:`~torch.isinf`
|
|
:func:`~torch.isposinf`
|
|
:func:`~torch.isneginf`
|
|
:func:`~torch.isnan`
|
|
:func:`~torch.erf`
|
|
:func:`~torch.erfinv`
|