Closed zimonitrome closed 2 years ago
Hi! What I would recommend is reshaping the tensor to be two dimensions, e.g [B, H, W] -> [B * H, W] and then expanding it after the operation. This is how some of the torch functions operate behind the scenes as well. This will be much faster than using a for loop
Here is some example code:
In [1]: import torch
In [2]: import torchsort
In [3]: B, H, W = 16, 10, 15
In [4]: x = torch.randn(B, H, W)
In [5]: x = x.view(B * H, W)
In [6]: x = torchsort.soft_sort(x)
In [7]: x = x.view(B, H, W)
In [8]: x.shape
Out[8]: torch.Size([16, 10, 15])
This is clever and I actually thoight of this after making the issue too. Will try it out, thanks!
Edit: For anyone interested, I was able to speed up my sorting by ~30x by calling soft_sort
only once.
I love this library now that I got it to work with my code!
I was wondering though, are there any plans to make it work with ambiguously shaped tensors? What work would that entail?
My current training scheme has tensors shaped [B, H, W], so I am currently doing
torch.stack([soft_sort(item) for item in batch])
. This library is fast, but running it sequentially like that is not. Maybe there is a way to parallelize it or extend the function to use more dimensions?