2

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!

ZZZ
  • 21
  • 4

2 Answers2

1

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]])
Shai
  • 111,146
  • 38
  • 238
  • 371
0

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.

ZZZ
  • 21
  • 4
  • have you cjecked if `torch.einsum` works with sparse tensors? – Shai Jan 14 '21 at 00:08
  • 1
    I just checked, it does not work with sparse tensors. But if I make them dense before, then it works (like tensordot). – ZZZ Jan 14 '21 at 09:21
  • moving to 6 dims: it is always going to be easier using `einsum` or `tensordot` however, since both functions do not support sparse tensors (yet), I think there's no reall alternative but the one outlined in my answer: reshaping to 2D and applying sparse dot product. Generalizing to 6 dim is not that complicated: you only need `coef` to be `2x6` instead of `2x4`. – Shai Jan 17 '21 at 08:14