TORCH.SPARSE稀疏存储格式学习笔记

介绍 

PyTorch对于稀疏存储格式的支持,目前主要支持COO和CSR格式。

目录

介绍 

Sparse COO Tensors

构造

Hybrid sparse COO tensors

Uncoalesced sparse COO tensors未合并稀疏COO张量

Working with sparse COO tensors

Sparse CSR tensors

---

Sparse COO Tensors

 PyTorch将坐标格式Coordinate Format（COO）存为元组tuple，一个元组包含一组索引indices（row, col）和对应的值(values).具体来讲，

• 索引indices长宽为tensor的维度和非零元素的数量（ndim，nse），元素类型为torch.int64。
• 对应的值values的长度为非零元素的数量（nse，）

COO格式的存储一般可以省掉200倍的存储空间占用。

构造

indice-i [[row, col]]

values-v

>>> i = [[1,3,0,2,4],
         [1,4,2,3,3]]
>>> v = [2,5,9,1,6]
>>> s = torch.sparse_coo_tensor(i, v, (6, 7))
>>> s.to_dense()

tensor([[0, 0, 9, 0, 0, 0, 0],
        [0, 2, 0, 0, 0, 0, 0],
        [0, 0, 0, 1, 0, 0, 0],
        [0, 0, 0, 0, 5, 0, 0],
        [0, 0, 0, 6, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 0]])

需要注意的是i不是一个索引元组的list，所以如果想要修改你的索引，你应该在实现稀疏构造器之前变换。

Hybrid sparse COO tensors

PyTorch支持清晰分开的（密集矩阵维度接在稀疏维度之后）混合稀疏COO存储格式，其中可以包含dense的部分。具体来讲，

• 索引indices长宽为tensor的维度和非零元素的数量（ndim，nse），元素类型为torch.int64。
• 对应的值values的长度为非零元素的数量（nse，dense_dims）

Uncoalesced sparse COO tensors未合并稀疏COO张量

可以在矩阵同一坐标索引上赋多个值。

>>> 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)

如果coalesce（）合并，即通过加和sum的方式。

>>> s.coalesce()
tensor(indices=tensor([[1]]),
       values=tensor([7]),
       size=(3,), nnz=1, layout=torch.sparse_coo)

Working with sparse COO tensors

查看是否是稀疏格式

>>> isinstance(s, torch.Tensor)
True
>>> s.is_sparse
True
>>> s.layout == torch.sparse_coo
True

查看稀疏和密集维度大小

>>> s.sparse_dim(), s.dense_dim()
(2, 1)

Coalesced COO的indices和values查看

>>> torch.Tensor.indices() 
>>> torch.Tensor.values().

Uncoalesced COO的indices查看

>>> s._indices()
tensor([[0, 1, 1],
        [2, 0, 2]])

切片操作slicing仅仅支持密集维度dense dimensions，指引indexing操作支持稀疏和密集维度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])

Sparse CSR tensors

类比COO