快捷方式

torch.sparse.mm

torch.sparse.mm()

执行稀疏矩阵 mat1 和(稀疏或带步长)矩阵 mat2 的矩阵乘法。类似于 torch.mm(),如果 mat1 是一个 (n×m)(n \times m) 张量,mat2 是一个 (m×p)(m \times p) 张量,输出将是一个 (n×p)(n \times p) 张量。当 mat1 是一个 COO 张量时,它必须具有 sparse_dim = 2。当输入是 COO 张量时,此函数还支持对两个输入的 backward 运算。

支持 CSR 和 COO 存储格式。

注意

此函数不支持计算相对于 CSR 矩阵的导数。

此函数还另外接受一个可选的 reduce 参数,该参数允许指定可选的约简操作,在数学上执行以下操作

zij=k=0K1xikykjz_{ij} = \bigoplus_{k = 0}^{K - 1} x_{ik} y_{kj}

其中 \bigoplus 定义了约简运算符。 reduce 仅在 CPU 设备上针对 CSR 存储格式实现。

参数
  • mat1 (张量) – 第一个要相乘的稀疏矩阵

  • mat2 (张量) – 第二个要相乘的矩阵,可以是稀疏的或稠密的

  • reduce (字符串, 可选) – 对非唯一索引应用的缩减操作 ("sum", "mean", "amax", "amin")。默认值 "sum"

形状

此函数的输出张量的格式如下: - 稀疏 x 稀疏 -> 稀疏 - 稀疏 x 稠密 -> 稠密

示例

>>> a = torch.tensor([[1., 0, 2], [0, 3, 0]]).to_sparse().requires_grad_()
>>> a
tensor(indices=tensor([[0, 0, 1],
                       [0, 2, 1]]),
       values=tensor([1., 2., 3.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo, requires_grad=True)
>>> b = torch.tensor([[0, 1.], [2, 0], [0, 0]], requires_grad=True)
>>> b
tensor([[0., 1.],
        [2., 0.],
        [0., 0.]], requires_grad=True)
>>> y = torch.sparse.mm(a, b)
>>> y
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseAddmmBackward0>)
>>> y.sum().backward()
>>> a.grad
tensor(indices=tensor([[0, 0, 1],
                       [0, 2, 1]]),
       values=tensor([1., 0., 2.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
>>> c = a.detach().to_sparse_csr()
>>> c
tensor(crow_indices=tensor([0, 2, 3]),
       col_indices=tensor([0, 2, 1]),
       values=tensor([1., 2., 3.]), size=(2, 3), nnz=3,
       layout=torch.sparse_csr)
>>> y1 = torch.sparse.mm(c, b, 'sum')
>>> y1
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)
>>> y2 = torch.sparse.mm(c, b, 'max')
>>> y2
tensor([[0., 1.],
        [6., 0.]], grad_fn=<SparseMmReduceImplBackward0>)

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得您的问题解答

查看资源