Exploiting sparsity in batch operations?

Is there a canonical way to exploit sparsity in batch operations torch.bmm() and torch.baddmm() yet? Interested mainly in sparse -> dense (dense -> sparse is also interesting).

If I have a batch of sparse input matrices, and a dense batch of matrices :

> mat1 = torch.zeros(4, 3, 5)
> mat1[1][1][1] = 1; mat1[2][2][2] = 1
> mat2 = torch.rand(4, 5, 6)
> torch.bmm(mat1, mat2)

Exploiting sparsity is quite an optimisation. If this isn’t available yet, which users might I liaise with to help out?

Edit: it seems there is a torch.SparseFloatTensor() available in a new release?

1 Like

bmm is currently not implemented for torch.sparse.* modules. If you can store it in a list, you can simply do torch.mm(mat1_i, mat2). If your matricies are extremely sparse, this should be pretty good :slight_smile:

Also, torch.smm(mat1_i, mat2) is also implemented, for sparse * dense -> sparse operations.

1 Like

Thanks for the tips :slight_smile: If my sparse matrices are in a list, do you mean something like this? :

import torch
import torch.autograd as ann

mat2 = ann.Variable(torch.rand(4, 5, 6), requires_grad=True)
mats = [ann.Variable(torch.zeros(4, 3, 5), requires_grad=True) for _ in range(3)]
for i in range(len(mats)):
    result = torch.bmm(mats[i], mat2)
    print result.size()

Something like this:

import torch

x = torch.rand(5,6)
# Sparse matrix of (0, 1) = 1; (2, 1) = 2, (3, 4) = 3
sparse = torch.sparse.FloatTensor(
    torch.LongTensor([[0, 2, 3], [1, 1, 4]]), # Indicies
    torch.FloatTensor([1, 2, 3])) # values

print(torch.mm(sparse, x))
# This won't actually save space or compute, since it's so dense,
# but it will be a sparse tensor representation.
print(torch.smm(sparse, x))

Simply construct a list of your sparse tensors, and loop over them to do the batch mm.

1 Like

Thanks once more! One last thing:

sparse = torch.sparse.FloatTensor(
    torch.LongTensor([[0, 2, 3], [1, 1, 4]]),   # Indicies
    torch.FloatTensor([1, 2, 3]))   # values

Seems strange to me as you don’t define the sizes of the sparse matrix - it seems to arbitrarily pick the indices of the corner value as the size. What is the logic here?

Check out the tests for more in depth use cases: https://github.com/pytorch/pytorch/blob/master/test/test_sparse.py

You can pass in a third argument to specify the size, like torch.sparse.FloatTensor(indicies, values, torch.Size([4, 5]))

Excellent @ebetica, thanks!

By the way, is anyone working on wrapping cusparse?

1 Like

martin raison is working on it, are you interested in collaborating?

Yes, sure! I’m not sure whether I’m familiar with pytorch internals enough to be able help. But I can try anyway.

I’d like to collaborate on writing a wrapper for cusparse, if you folks still need a hand.

@siddharthachandra have a look at https://github.com/pytorch/pytorch/pull/1147
Part of it is done.

Thanks, looks good. Let me check it out.