将 tensordot 与 torch.sparse 张量一起使用

Question

是否可以使用与“tensordot”类似的方法来处理torch.sparse 张量？

我正在尝试将 4 维张量应用于 2 维张量。这可以使用 torch 或 numpy 来实现。但是，在不使用“.to_dense()”使稀疏张量密集的情况下，我没有找到使用 torch.sparse 的方法。

更准确地说，这是我想要在不使用“.to_dense()”的情况下做的事情：

import torch
import torch.sparse

nb_x = 4
nb_y = 3
coordinates = torch.LongTensor([[0,1,2],[0,1,2],[0,1,2],[0,1,2]])
values = torch.FloatTensor([1,2,3])
tensor4D = torch.sparse.FloatTensor(coordinates,values,torch.Size([nb_x,nb_y,nb_x,nb_y]))
inp = torch.rand((nb_x,nb_y))

#what I want to do
out = torch.tensordot(tensor4D.to_dense(),inp,dims=([2,3],[0,1]))

print(inp)
print(out)

（这里是输出：torch_code）

或者，这里是使用 numpy 的类似代码：

import numpy as np

tensor4D = np.zeros((4,3,4,3))
tensor4D[0,0,0,0] = 1
tensor4D[1,1,1,1] = 2
tensor4D[2,2,2,2] = 3
inp = np.random.rand(4,3)

out = np.tensordot(tensor4D,inp)

print(inp)
print(out)

（这里是输出：numpy_code）

感谢您的帮助！

Answer 1

您的特定tensordot 可以通过“挤压”tensor4D 的前两个和后两个维度转换为简单的矩阵乘法。

简而言之，你想做的是

raw = tensor4D.view(nb_x*nb_y, nb_x*nb_y) @ inp.flatten()
out = raw.view(nb_x, nb_y)

但是，由于 view 和 reshape 未针对稀疏张量实现，因此您必须手动执行：

sz = tensor4D.shape
coeff = torch.tensor([[1, sz[1], 0, 0], [0, 0, 1, sz[3]]])
reshaped = torch.sparse.FloatTensor(coeff @ idx, tensor4D._values(), torch.Size([nb_x*nb_y, nb_x*nb_y]))

# once we reshaped tensord4D it's all downhill from here
raw = torch.sparse.mm(reshaped, inp.flatten()[:, None])
out = raw.reshape(nb_x, nb_y)
print(out)

输出是

tensor([[0.4180, 0.0000, 0.0000],
   [0.0000, 0.6025, 0.0000],
   [0.0000, 0.0000, 0.5897],
   [0.0000, 0.0000, 0.0000]])

Answer 2

确实，这个效果很好，谢谢你的回答！

这种方法的弱点在我看来很难一概而论。

实际上，“inp”和“out”应该是图像。在这里，它们是黑白图像，因为只有两个维度：高度和宽度。

如果相反，我拍摄 RGB 图像，那么我将不得不考虑作用于 3D 张量的 6D 张量。我仍然可以通过将前三个维度和后三个维度“挤压”在一起来应用相同的技巧。但是在我看来，它会很快变得更加参与（也许我错了）。而使用 tensordot 来泛化会更简单。

因此，我将使用您提出的解决方案，但如果有人找到其他解决方案，我仍然很感兴趣。