pytorch稀疏张量sparse tensor中关于coalesce()的用法

定义一个稀疏tensor 其中可能出现同一个索引对应多个标量

sp_adj torch.sparse_coo_tensor([[1,1,1,2]], [2,4,6,9], (3,))
print(sp_adj)
output:tensor(indices tensor([[1, 1, 1, 2]]),
 values tensor([2, 4, 6, 9]),
 size (3,), nnz 4, layout torch.sparse_coo)

coalesce()函数的作用就是对相同索引的多个值求和

a a.coalesce()
print(a)
output:tensor(indices tensor([[1, 2]]),
 values tensor([12, 9]),
 size (3,), nnz 2, layout torch.sparse_coo)

官方文档如下

另外值得注意的是 pytorch中关于两个稀疏tensor的求和运算只是简单地对各自的indices和values进行叠加

a torch.sparse_coo_tensor([[1,1,2,2]], [2,4,6,9], (3,))
b torch.sparse_coo_tensor([[0,0,0,2]], [1,3,5,7], (3,))
c a b
print(c)
output:tensor(indices tensor([[0, 0, 0, 1, 1, 2, 2]]),
 values tensor([ 1, 3, 5, 2, 4, 13, 9]),
 size (3,), nnz 7, layout torch.sparse_coo)

因此多次执行稀疏tensor的加法运算可能导致结果变得越来越大 这在官方文档也有说明 因此再执行完加法运算后 使用.coalesce()进行压缩