CODEWITHSUNDEEP
pytorchsparse-matrixtensortensordot
ZZZ
ZZZ

Reputation: 61

Using tensordot with torch.sparse tensors

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!

Upvotes: 5

Views: 1651

Answers (3)

user115168
user115168

Reputation: 1

In another answer, a reshaping coeff matrix is proposed for shaping the indices of the 4D matrix:

raw = tensor4D.view(nb_x*nb_y, nb_x*nb_y) @ inp.flatten()
out = raw.view(nb_x, nb_y)
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]))
raw = torch.sparse.mm(reshaped, inp.flatten()[:, None])
out = raw.reshape(nb_x, nb_y)

I was able to get the same answer as tensordot if I change the coeff to be instead:

coeff = torch.tensor([[sz[1], 1, 0, 0], [0, 0, sz[3], 1]])

This is because flattening a 2D (x, y) coordinate to 1D would be x*dim(y) + y rather than x + dim(y)*y.

(This is a comment on Shai's answer but since I'm starting a new account I would have to start an independent answer).

Upvotes: 0

ZZZ
ZZZ

Reputation: 61

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.

Upvotes: 1

Shai
Shai

Reputation: 114796

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

Upvotes: 1

Related Questions