soft_sort and soft_rank don't work for pytorch cuda tensors about fast-soft-sort HOT 4 CLOSED

So I got this more complete example to work where I apply the soft_sort operator on a parameter of a nn.Module on a cuda device. I don't know how general my approach is and whether it reliably solves the issue for all scenarios but I figured I would share it anyway:

import torch
import torch.nn as nn
from fast_soft_sort import pytorch_ops


class Net(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.tensor([[0.1, 0.7, 0.2]]))

    def forward(self, x):
        w_sorted = pytorch_ops.soft_sort(self.weight.cpu()).cuda()
        return 2.0 * w_sorted * x


net = Net()
net = net.cuda()

x = torch.tensor([[4.0, 5.0, 6.0]]).cuda()
y = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float64).cuda()

yhat = net(x)
loss = nn.functional.mse_loss(yhat, y)
loss.backward()

print("Loss = {:.2f}".format(loss))
print(f"Grad = {net.weight.grad}")

Still curious to see whether there is a better way to handle that internally.

from fast-soft-sort.

At the moment we do not have a GPU implementation of the projection operators, which is the cause for the error. We decided not to do this conversion implicitly as we want the user to be aware that a device copy is necessary. If you want that behavior, can you write a small util function like

def soft_sort(array):
   return pytorch_ops.soft_sort(array.cpu()).cuda()

and then use it as a plug-in replacement. Would that work?

from fast-soft-sort.

Thank you for your help! Yes, it does work indeed. :)

from fast-soft-sort.

I have been working on a pure PyTorch implementation here: https://github.com/teddykoker/torchsort, complete with the isotonic regression code written as a C++ extension. The CPU implementation is much faster (see benchmarks), and I am working on the CUDA implementation which should be done soon.

from fast-soft-sort.