PyTorch 扩展PyTorch Sparse

PyTorch Sparse

pytorch_sparse2026-08-24 19:46:25

PyTorch Sparse 是一个用于稀疏张量操作的小型 PyTorch 扩展库，该扩展库包含具有 autograd 支持的优化稀疏矩阵运算。

仓库
https://github.com/rusty1s/pytorch_sparse

PyPI
https://pypi.org/project/torch-sparse/

安装

pip install torch-sparse

基本使用

稀疏矩阵与密集矩阵的乘法 (Sparse Dense Matrix Multiplication)

spmm 函数用于稀疏矩阵和密集矩阵的乘积。

torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor

参数

  • index（LongTensor） - 稀疏矩阵的索引张量。
  • value（Tensor） - 稀疏矩阵的值张量。
  • m（int） - 稀疏矩阵的第一个维度。
  • n（int） - 稀疏矩阵的第二个维度。
  • matrix（Tensor） - 密集矩阵。

返回

  • out（Tensor） - 密集矩阵。

示例

import torch
from torch_sparse import spmm

index = torch.tensor([[0, 0, 1, 2, 2],
                      [0, 2, 1, 0, 1]])
value = torch.Tensor([1, 2, 4, 1, 3])
matrix = torch.Tensor([[1, 4], [2, 5], [3, 6]])

out = spmm(index, value, 3, 3, matrix)

print(out)
# tensor([[7.0, 16.0],
#         [8.0, 20.0],
#         [7.0, 19.0]])

稀疏矩阵之间的乘法 (Sparse Sparse Matrix Multiplication)

spspmm 函数用于计算两个稀疏矩阵的乘积。

torch_sparse.spspmm(indexA, valueA, indexB, valueB, m, k, n) -> (torch.LongTensor, torch.Tensor)

参数

  • indexA（LongTensor） - 第一个稀疏矩阵的索引张量。
  • valueA（Tensor） - 第一个稀疏矩阵的值张量。
  • indexB（LongTensor） - 第二个稀疏矩阵的索引张量。
  • valueB（Tensor） - 第二个稀疏矩阵的值张量。
  • m（int） - 第一个稀疏矩阵的第一个维度。
  • k（int） - 第一个稀疏矩阵的第二个维度和第二个稀疏矩阵的第一个维度。
  • n（int） - 第二个稀疏矩阵的第二个维度。
  • coalesced（bool， optional）：如果设置为 True，将合并两个输入稀疏矩阵。（默认值：False）

返回

  • index（LongTensor） - 稀疏矩阵的输出索引张量。
  • value（Tensor） - 稀疏矩阵的输出值张量。

示例

import torch
from torch_sparse import spspmm

indexA = torch.tensor([[0, 0, 1, 2, 2], [1, 2, 0, 0, 1]])
valueA = torch.Tensor([1, 2, 3, 4, 5])

indexB = torch.tensor([[0, 2], [1, 0]])
valueB = torch.Tensor([2, 4])

indexC, valueC = spspmm(indexA, valueA, indexB, valueB, 3, 3, 2)

print(indexC)
# tensor([[0, 1, 2],
#         [0, 1, 1]])
print(valueC)
# tensor([8.0, 6.0, 8.0])

本页目录