pytorch笔记：torch.sparse类

UQI-LIUWJ

已于 2022-10-20 10:27:57 修改

阅读量1w

点赞数 13

分类专栏： pytorch学习

于 2021-10-10 15:18:34 首次发布

本文链接：https://blog.csdn.net/qq_40206371/article/details/120686280

版权

PyTorch 稀疏张量 COO格式矩阵运算内存效率

关键词由CSDN通过智能技术生成

pytorch学习专栏收录该内容

116 篇文章

订阅专栏

PyTorch 提供了 torch.Tensor 来表示一个包含单一数据类型元素的多维数组。默认情况下，数组元素连续存储在内存中，从而可以有效地实现各种数组处理算法，这些算法依赖于对数组元素的快速访问。

然而，存在一类重要的多维数组，即所谓的稀疏数组，其中数组元素的连续内存存储被证明是次优的。 稀疏数组具有大部分元素为零的特性，这意味着如果仅存储或/和处理非零元素，则可以节省大量内存和处理器资源。

1 构造稀疏矩阵

import torch
i = torch.LongTensor([[0, 1, 1],[2, 0, 2]])   #row, col
v = torch.FloatTensor([3, 4, 5])    #data
torch.sparse.FloatTensor(i, v, torch.Size([2,3])).to_dense()   #torch.Size
'''
tensor([[0., 0., 3.],
        [4., 0., 5.]])
'''

构造方法和 scipy笔记：scipy.sparse_UQI-LIUWJ的博客-CSDN博客 2.2 coo矩阵的类似

在 PyTorch 中，稀疏张量的填充值不能明确指定，一般假定为零。
但是，存在可能以不同方式解释填充值的操作。例如，torch.sparse.softmax() 计算 softmax 时假设填充值为负无穷大。

1.1 稀疏COO tensor

在 COO 格式（coordinate）中，指定的元素存储为元素索引和相应值的元组。
- 元素索引的类型是torch.int64，size是(ndim,nse)
- 数值类型是任何类型，size是(nse,)

import torch

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.is_sparse
#True

s.layout
#torch.sparse_coo

x坐标为0，y坐标为2的元素是3；x坐标为1，y坐标为1的元素是4.。。。

1.1.0 稀疏矩阵转换成正常Tensor

s.to_dense()
'''
tensor([[0, 0, 3],
        [4, 0, 5]])
'''

1.1.1 COO tensor和正常tensor（strided tensor）空间复杂度对比

COO tensor：ndim*8*nse+<element_size>*nse
正常tensor: <tensor_size>*<element_size>

举例一个10,000*10,000的 float32 Tensor，其中有100,000个非零元素
- 正常Tensor的话，需要10 000 * 10 000 * 4 = 400 000 000 比特
- 使用COO tensor的话，需要(2 * 8 + 4) * 100 000 = 2 000 000比特

1.2.3 创建空的COO tensor

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

1.2.4 混合稀疏COO Tensor

我们可以将前面value值为标量的稀疏张量扩展到value值为连续张量的稀疏张量。
这种张量称为混合张量。
- 元素索引的类型是torch.int64，size是(ndim,nse)
- 数值类型是任何类型，size是(nse,dense_dims)
对应的稀疏矩阵的维度是n_dim+dense_dims

import torch

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [[3,2],[4,1],[5,3]]

s = torch.sparse_coo_tensor(i, v, (2, 3,2))
s,s.to_dense()

'''
(tensor(indices=tensor([[0, 1, 1],
                        [2, 0, 2]]),
        values=tensor([[3, 2],
                       [4, 1],
                       [5, 3]]),
        size=(2, 3, 2), nnz=3, layout=torch.sparse_coo),
 tensor([[[0, 0],
          [0, 0],
          [3, 2]],
 
         [[4, 1],
          [0, 0],
          [5, 3]]]))
'''

1.2.5 未合并的稀疏COO 张量

PyTorch 稀疏 COO 张量格式允许未合并的稀疏张量，其中索引中可能存在重复坐标；
在这种情况下，该索引处的值是所有重复值条目的总和。
例如，可以为同一个索引 1 指定多个值 3 和 4，这会导致未合并张量：

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)
'''
s.is_coalesced()
#False

合并（结果仍为稀疏张量）

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

s.to_dense()
#tensor([0, 7, 0])

1.2.6 是否需要合并？

在大多数情况下，不用关心稀疏张量是否被合并，因为在给定合并或未合并稀疏张量的情况下，大多数操作的工作方式相同。
但是，一些操作可以在未合并的张量上更有效地实现，而一些操作可以在合并的张量上更有效地实现。
例如，通过简单地连接索引和值张量来实现稀疏 COO 张量的添加：

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

1.2.7 查看indice和value

1）不用事先合并

import torch

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [3, 4, 5]

s = torch.sparse_coo_tensor(i, v, (2, 3))
print(s)
print(s._indices())
print(s._values())

'''
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
tensor([[0, 1, 1],
        [2, 0, 2]])
tensor([3, 4, 5])
'''

2) 需要事先合并

print(s.indices())
print(s.values())

'''
RuntimeError                              Traceback (most recent call last)
<ipython-input-27-b4753553cd54> in <module>
      8 s = torch.sparse_coo_tensor(i, v, (2, 3))
      9 print(s)
---> 10 print(s.indices())
     11 print(s.values())

RuntimeError: Cannot get indices on an uncoalesced tensor, please call .coalesce() first
'''

print(s.coalesce().indices())
print(s.coalesce().values())
'''
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
tensor([[0, 1, 1],
        [2, 0, 2]])
tensor([3, 4, 5])
'''

1.2.6 sparse_dim 和dense_dim

一个是index的dim，一个是value的dim

import torch

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [3, 4, 5]

s = torch.sparse_coo_tensor(i, v, (2, 3))
print(s.sparse_dim(),s.dense_dim())
#（2,0）

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [[3,2],[4,1],[5,3]]

s = torch.sparse_coo_tensor(i, v, (2, 3,2))
print(s.sparse_dim(),s.dense_dim())
#(2,1)

1.2.7 切片和索引

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [[3,2],[4,1],[5,3]]

s = torch.sparse_coo_tensor(i, v, (2, 3,2))
print(s)
'''
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([[3, 2],
                      [4, 1],
                      [5, 3]]),
       size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)
'''

在sparse维度（index部分）和dense部分（value部分）都可以索引

s[1]
'''
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([[3, 2],
                      [4, 1],
                      [5, 3]]),
       size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)
'''

s[1,0,1]
#tensor(1)

切片只能在dense部分切

s[1,0,1:],s[1,0,0:]
#(tensor([1]), tensor([4, 1]))

2 稀疏矩阵的基本运算

基本上都是第一个参数是sparse的，第二个是正常Tensor

先构造两个稀疏矩阵

import torch
i = torch.LongTensor([[0, 1, 1],[2, 0, 2]])   #row, col
v = torch.FloatTensor([3, 4, 5])    #data
x1=torch.sparse.FloatTensor(i, v, torch.Size([2,3]))

x1,x1.to_dense()  
'''
(tensor(indices=tensor([[0, 1, 1],
                        [2, 0, 2]]),
        values=tensor([3., 4., 5.]),
        size=(2, 3), nnz=3, layout=torch.sparse_coo),
 tensor([[0., 0., 3.],
         [4., 0., 5.]]))
'''

import torch
i = torch.LongTensor([[0, 1, 1],[1, 0, 1]])   #row, col
v = torch.FloatTensor([3, 4, 5])    #data
x2=torch.sparse.FloatTensor(i, v, torch.Size([3,2]))

x2,x2.to_dense() 
'''
(tensor(indices=tensor([[0, 1, 1],
                        [1, 0, 1]]),
        values=tensor([3., 4., 5.]),
        size=(3, 2), nnz=3, layout=torch.sparse_coo),
 tensor([[0., 3.],
         [4., 5.],
         [0., 0.]]))
'''

2.0 不支持dense * sparse！

pytorch不支持M[strided] @ M[sparse_coo]

如果需要，可以这么整：D @ S == (S.t() @ D.t()).t()

2.1 稀疏矩阵的乘法

2.1.1 torch.mm

只支持第二个参数是dense（即dense*dense，或者sparse*dense），输出是dense

dense*dense
dense*sparse
sparse*sparse
sparse*dense

2.1.2 torch.sparse.mm

同样地，只支持第二个参数是dense（即dense*dense，或者sparse*dense）

dense*dense
dense*sparse
sparse*sparse
sparse*dense

2.1.3 torch.mv

矩阵和向量的乘法，第二个也只能是dense的

import torch

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [3, 4, 5]

s = torch.sparse_coo_tensor(i, v, (2, 3))
print(s.to_dense())
'''
tensor([[0, 0, 3],
        [4, 0, 5]])
'''


t=torch.LongTensor([1,2,3])
torch.mv(s,t),s@t
'''
(tensor([ 9, 19]), tensor([ 9, 19]))
'''

2.1.4 torch.matmul

和torch.mm 类似，第二个也是只能dense

import torch

i = [[0, 1, 1],
    [2, 0, 2]]

v =  [3, 4, 5]

s = torch.sparse_coo_tensor(i, v, (2, 3))

t=torch.LongTensor([[1],[2],[3]])
torch.matmul(s,t),s@t

'''
(tensor([[ 9],
         [19]]),
 tensor([[ 9],
         [19]]))
'''

2.2 转置

t()即可

x2,x2.to_dense()
'''
(tensor(indices=tensor([[0, 1, 1],
                        [1, 0, 1]]),
        values=tensor([3., 4., 5.]),
        size=(3, 2), nnz=3, layout=torch.sparse_coo),
 tensor([[0., 3.],
         [4., 5.],
         [0., 0.]]))
'''

x2.t(),x2.t().to_dense()

'''
(tensor(indices=tensor([[1, 0, 1],
                        [0, 1, 1]]),
        values=tensor([3., 4., 5.]),
        size=(2, 3), nnz=3, layout=torch.sparse_coo),
 tensor([[0., 4., 0.],
         [3., 5., 0.]]))
'''

2.3 索引

稀疏矩阵支持整行索引，支持Sparse.matrix[row_index]；

x2,x2.to_dense()
'''
(tensor(indices=tensor([[0, 1, 1],
                        [1, 0, 1]]),
        values=tensor([3., 4., 5.]),
        size=(3, 2), nnz=3, layout=torch.sparse_coo),
 tensor([[0., 3.],
         [4., 5.],
         [0., 0.]]))
'''

x2[1],x2[1].to_dense()
'''
(tensor(indices=tensor([[0, 1]]),
        values=tensor([4., 5.]),
        size=(2,), nnz=2, layout=torch.sparse_coo),
 tensor([4., 5.]))
'''

稀疏矩阵不支持具体位置位置索引Sparse.matrix[row_index,col_index]

x2[1][1],x2[1][1].to_dense()

2.4 相加

a = torch.sparse.FloatTensor(
    torch.tensor([[0,1,2],[2,3,4]]), 
    torch.tensor([1,1,1]), 
    torch.Size([5,5]))
a.to_dense()
'''
tensor([[0, 0, 1, 0, 0],
        [0, 0, 0, 1, 0],
        [0, 0, 0, 0, 1],
        [0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0]])
'''

a1=torch.sparse.FloatTensor(
    torch.tensor([[0,3,2],[2,3,2]]), 
    torch.tensor([1,1,1]), 
    torch.Size([5,5]))
a1.to_dense()
'''
tensor([[0, 0, 1, 0, 0],
        [0, 0, 0, 0, 0],
        [0, 0, 1, 0, 0],
        [0, 0, 0, 1, 0],
        [0, 0, 0, 0, 0]])
'''

只支持sparse+sparse

torch.add(a,a1) ,torch.add(a,a1).to_dense()
'''
(tensor(indices=tensor([[0, 1, 2, 3, 2],
                        [2, 3, 4, 3, 2]]),
        values=tensor([2, 1, 1, 1, 1]),
        size=(5, 5), nnz=5, layout=torch.sparse_coo),
 tensor([[0, 0, 2, 0, 0],
         [0, 0, 0, 1, 0],
         [0, 0, 1, 0, 1],
         [0, 0, 0, 1, 0],
         [0, 0, 0, 0, 0]]))
'''

a.add(a1),a.add(a1).to_dense()
#同理