Pytorch-稀疏张量 | 代码大全

Pytorch sparse tensors API目前还处于beta版本, 未来有变更的可能性!

默认情况下Pytorch将torch.Tensor中的元素连续地存储到物理存储中，这样有利于高效地实现数组操作的算法。但是对于某些高度稀疏的场景却不太适用了，比如图邻接矩阵中的大部分元素都是0，Pytorch通过稀疏存储格式(sparse storage formats)来提供更好的性能优化支持。

PyTorch目前支持COO、CSR、CSC、BSR、BSC格式，它们都是压缩数据来高效表示0值元素。

COO格式

COO格式中，非0元素通过tuple的形式进行存储，tuple由indices和对应的值组成。indices是一个二维数组，第一个元素是非0元素的行坐标，第二个元素是非0元素的列坐标。

对于下面这样的一个矩阵可以通过下面的代码来创建
[[0, 0, 3],
[4, 0, 5]]

i = [[0, 1, 1], [2, 0, 2]]
v =  [3, 4, 5]
s = torch.sparse_coo_tensor(i, v, (2, 3))

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

实际应用

Wide & Deep是一个比较经典的深度模型，其中wide部分是人工特征交叉组合，比如将机型和商品进行组合生成新的特征，而这种特征非常非常稀疏，正好适用于这种场景。

需要注意的一个地方是稀疏向量的维度，还是以Wide & Deep举例，假设wide模型有10个特征，每个特征有100维，那么wide的linear模型的维度应该是1000维。此时coo tensor的取值需要是0-999，而不是每个列的取值是0-99，因为coo会推断能够存储其中最大的值就可以了。

sparse_coo_tensor的初始化函数中有个size参数，官网文档是这样说的:
size (list, tuple, or torch.Size, optional) – Size of the sparse tensor. If not provided the size will be inferred as the minimum size big enough to hold all non-zero elements.