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?

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.

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(x)
print(sparse.to_dense())
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.

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?

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.