Using tensordot with torch.sparse tensors

user_14998703 0
Viewed 709

Is it possible to use a similar method as "tensordot" with torch.sparse tensors?

I am trying to apply a 4 dimensional tensor onto a 2 dimensional tensor. This is possible using torch or numpy. However, I did not find the way to do it using torch.sparse without making the sparse tensor dense using ".to_dense()".

More precisely, here is what I want to do without using ".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)

(here is the output: torch_code)

Alternatively, here is a similar code using 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)

(here is the output: numpy_code)

Thanks for helping!

tensordot torch tensor pytorch sparse-matrix
2 Answers

Your specific tensordot can be cast to a simple matrix multiplication by "squeezing" the first two and last two dimensions of tensor4D.

In short, what you want to do is

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

However, since view and reshape are not implemented for sparse tensors, you'll have to it manually:

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)

And the output is

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]])
user_17144100

Indeed, this works very well, thank you for your answer!

The weakness of this method seems to me that it is hard to generalize.

In fact, "inp" and "out" are supposed to be images. Here, they are black and white images since there are only two dimensions: height and width.

If instead, I take RGB images, then I will have to consider 6D tensors acting on 3D tensors. I can still apply the same trick by "squeezing" the first three dimensions together and the last three dimensions together. However it seems to me that it will become more involving very quickly (maybe I am wrong). While using tensordot instead would be much more simpler for generalization.

Therefore, I am going to use the solution you proposed, but I am still interested if someone finds an other solution.

user_149987030
Related