PyTorch: Zero all elements of vector except top k?

user_13456947 0
Viewed 811

I am trying to create a new activation layer, let’s call it topk, that would work as follows. It will take a vector x of size n as input (result of multiplying previous layer output by weight matrix and adding bias) and a positive integer k and would output a vector topk(x) of size n whose elements are:

              x_i (if x_i is one of the top k elements of x) 
topk(x)_i = 
              0 (otherwise)

While calculating gradient of topk(x), top k elements of x should have gradient 1, everything else 0.

How should I implement this? Any help will be appreciated.

top-n pytorch python
1 Answers

You can use torch.topk for this:

k = 2
output = torch.randn(5)
vals, idx = output.topk(k)

topk = torch.zeros_like(output)
topk[idx] = vals

>>> topk
tensor([1.0557, 0.0000, 0.0000, 1.4562, 0.0000])

Note that while the 'values' of topk() are differentiable, the 'indices' are not (similar to how argmax is not a differentiable function).

user_90676150
Related