Python PyTorch mm用法及代碼示例

本文簡要介紹python語言中 torch.sparse.mm 的用法。

用法:

torch.sparse.mm(mat1, mat2)

參數：

• mat1(SparseTensor) -第一個要相乘的稀疏矩陣

• mat2(Tensor) -要相乘的第二個矩陣，可以是稀疏的或密集的

執行稀疏矩陣 mat1 和(稀疏或跨步)矩陣 mat2 的矩陣乘法。與 torch.mm() 類似，如果 mat1 是 (n \times m) 張量，mat2 是 (m \times p) 張量，則 out 將是 (n \times p) 張量。 mat1 需要有 sparse_dim = 2 。此函數還支持兩個矩陣的後向。請注意，mat1 的梯度是一個合並的稀疏張量。

形狀：

該函數的輸出張量格式如下： - 稀疏 x 稀疏 -> 稀疏 - 稀疏 x 密集 -> 密集

例子：

>>> a = torch.randn(2, 3).to_sparse().requires_grad_(True)
>>> a
tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
                       [0, 1, 2, 0, 1, 2]]),
       values=tensor([ 1.5901,  0.0183, -0.6146,  1.8061, -0.0112,  0.6302]),
       size=(2, 3), nnz=6, layout=torch.sparse_coo, requires_grad=True)

>>> b = torch.randn(3, 2, requires_grad=True)
>>> b
tensor([[-0.6479,  0.7874],
        [-1.2056,  0.5641],
        [-1.1716, -0.9923]], requires_grad=True)

>>> y = torch.sparse.mm(a, b)
>>> y
tensor([[-0.3323,  1.8723],
        [-1.8951,  0.7904]], grad_fn=<SparseAddmmBackward>)
>>> y.sum().backward()
>>> a.grad
tensor(indices=tensor([[0, 0, 0, 1, 1, 1],
                       [0, 1, 2, 0, 1, 2]]),
       values=tensor([ 0.1394, -0.6415, -2.1639,  0.1394, -0.6415, -2.1639]),
       size=(2, 3), nnz=6, layout=torch.sparse_coo)