稀疏张量乘法

Question

我在 ytorch 工作，正在寻求有关张量乘法的指导。我拥有形状为 [b, n, n, i] 和 [i, k] 的张量。我的目标是执行乘法并获得形状为 [b, n, n, k] 的张量。第一个张量是稀疏的，我有兴趣利用它的稀疏性。是否有一个库或包支持此任务的高效稀疏张量乘法？

我尝试了 torch.sparse 但它仅支持二维张量。

Answer 1

您可以使用

torch.matmul

沿最后一个维度进行矩阵乘法。

import torch
import torch.nn as nn

x = torch.randn(12, 8, 8, 4)
y = torch.randn(4, 16)

z = torch.matmul(x,y)
z.shape
>torch.Size([12, 8, 8, 16])

关于稀疏性，您可以查看torch.sparse，但我认为除非两个张量都是稀疏的，否则它不会显着改善事情。