Search code examples
pythondeep-learningpytorchcoremlcoremltools

How to replace torch.sparse with other pytorch function?


I would like to replace the torch.sparse function using the other Pytorch function.

i = torch.LongTensor([[0, 1, 1], [2, 0, 2]])
v = torch.FloatTensor([3, 4, 5])
out1 = torch.sparse.FloatTensor(i, v, torch.Size([2,3])).to_dense()
out2 = ???
out1 == out2
> tensor(True)

Background:

I'm converting a Pytorch model to CoreML, but the sparse_coo_tensor operator defined in the torch.norm function is not implemented with CoreMLTools.

A few people have trouble with this problem, but CoreMLTools is still unsupported. So I'd like to replace it without the operator used in torch.sparse.FloatTensor.

I have tried torch.sparse_coo_tensor but it was not supported.

I have created a simple colaboratory notebook that reproduces this. Please test it using the following colab. https://colab.research.google.com/drive/1TzpeJqEcmCy4IuXhhl6LChFocfZVaR1Q?usp=sharing

I've asked about different operators on stackoverflow before, so please refer to that. How to replace torch.norm with other pytorch function?


Solution

  • If I understand sparse_coo format correctly, the two rows of i are just the coordinates at which to copy the values of v. Which means that you can instead create you matrix like :

    def dense_from_coo(i, v):
        rows = i[0].max()+1
        cols = i[1].max()+1
        out = torch.zeros(rows, cols)
        out[i[0], i[1]] = v
        return out
    
    print(dense_from_coo(i,v))
    >>> tensor([[0., 0., 3.],
                [4., 0., 5.]])
    
    print(torch.sparse.FloatTensor(i, v, torch.Size([2,3])).to_dense())
    >>> tensor([[0., 0., 3.],
                [4., 0., 5.]])