pytorch稀疏矩阵处理（torch.sparse）

Pytorch稀疏矩阵处理

本文将简单介绍稀疏矩阵常用的存储方式和Pytorch中稀疏矩阵的处理方法。常用的稀疏矩阵存储格式有COO，CSR/CSC，LIL。

稀疏矩阵存储方式

1. COO

COO(Coordinate format )是最为简单的格式，以三元组的形式存储稀疏矩阵。记录矩阵中非零元素的数值和所在的行序号和列序号。形式为（行号，列号，数值）。这种存储方式的主要优点是灵活、简单。但是缺点是不可以直接进行矩阵的相关运算。

2. CSR/CSC

CSR（Compressed Sparse Row）格式实现了用于存储二维张量的 CSR 格式。尽管不支持 N 维张量，但与 COO 格式相比的主要优势是更好地利用存储和更快的计算操作。目前尚不存在 CUDA 支持。

3. LIL

LIL (List-of-List) 每行存储一个列表，每个条目包含列索引和值。通常，这些条目按列索引进行排序，以便更快地查找。

稀疏矩阵的处理

Pytorch中，处理稀疏矩阵的有效工具torch.sparse。Torch 支持 COO(rdinate) 格式的稀疏张量，可以有效地存储和处理大多数元素为零的张量。下面主要介绍torch.sparse处理COO形式矩阵的方法。文中使用的Pytorch的版本为1.7.1

1.torch.sparse.FloatTensor类

使用该类，可定义一个COO类型的稀疏矩阵。

# sparse tensor
i = torch.LongTensor([[0, 1, 1],
                     [2, 1, 0]])
d = torch.tensor([3, 6, 9], dtype=torch.float)
a = torch.sparse.FloatTensor(i, d, torch.Size([2, 3]))
print(a)

得到的输出
tensor(indices=tensor([[0, 1, 1],
                       [2, 1, 0]]),
       values=tensor([3., 6., 9.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)

i 为非零元素的索引， d为非零元素的值。根据COO表示的规则，可以得出，在（0，2）位置为3，（1，1）为6，（1，0）位置为9。使用to_dense()方法可以将COO形式的矩阵转换为普通的Tensor形式，在定义时还需要指定原来稀疏矩阵的大小，使用torch.Size([2,3])作为参数定义原有稀疏矩阵的大小。该类还有其他的一些操作，比如add()或者sub()等

print（a.to_dense())
print(a.add(a).to_dense())  # a + a

输出为
tensor([[0., 0., 3.],
        [9., 6., 0.]])
        
tensor([[ 0.,  0.,  6.],
        [18., 12.,  0.]])

2.torch.sparse.mm

该函数的原型为torch.sparse.mm(mat1: torch.Tensor, mat2: torch.Tensor) → torch.Tensor。特别注意第一个矩阵（mat1）的数据类型为SparseTensor,第二个矩阵为Tensor类型。支持backward操作。

mat1 (SparseTensor) – the first sparse matrix to be multiplied
mat2 (Tensor) – the second dense matrix to be multiplied
return - torch.Tensor

b = torch.eye(3)  # 定义一个单位矩阵
y = torch.sparse.mm(a, b)

输出为
tensor([[0., 0., 3.],
        [9., 6., 0.]])

3.torch.sparse.sum

该函数的原型为

torch.sparse.sum（input：torch.Tensor，dim：Optional[Tuple[int]] = None，dtype：Optional[int] = None ） → torch.Tensor

返回值为指定维度的和，通过dim参数来指定维度，dim为一个Tuple

x = torch.sparse.sum(a, dim=[0])  # 对列求和
tensor([9., 6., 3.])
x = torch.sparse.sum(a, dim=[1])  # 对行求和
tensor([ 3., 15.])

参考资料

Pytorch官方文档（Version = 1.7.1）