How can I use argsort in Keras

I study ML and I using a open source(blstm) and I want adding a argsort layer, but keras don't have a argsort layer.

I searched this question, and people said that use a Lambda function. (I use Tensorflow backend) but I don't know that how can I defined function and how can i use that.

This is original code:

tagger = Dropout(self.dropout_ratio)(tagger)
prediction = TimeDistributed(Dense(self.output_vocab_size, 
activation='softmax'))(tagger)
self.model = Model(input=raw_current, output=prediction)
self.model.compile(loss='categorical_crossentropy', optimizer=opt_func)

prediction shape is (?,48,224).

After trainging, when I test this model, I want get a top_10 index from prediction[0][-1], so i use some sorting code.
this is my test code.

prediction = self.model.predict(test_batch_data)
my_prediction = prediction[0][-1]
top10_my_prediction_idx = sorted(range(len(rule_prediction)), key=lambda k: rule_prediction[k] , reverse=True)[0:10]

I want add a get_top_10 layer after prediction layer when I train model.

Like this one:

tagger = Dropout(self.dropout_ratio)(tagger)
prediction = TimeDistributed(Dense(self.output_vocab_size, activation='softmax'))(tagger)
**top10_prediction = Lambda(get_top10_prediction)**
self.model = Model(input=raw_current, output=prediction)
self.model.compile(loss='categorical_crossentropy', optimizer=opt_func)

how can I define Lambda function and how can I use that?

Solution

If you're using tensorflow backend, the tensor can be arg-sorted via tf.nn.top_k function with sorted=True argument. A lambda layer will look like this:

def top_k(input, k):
  # Can also use `.values` to return a sorted tensor
  return tf.nn.top_k(input, k=k, sorted=True).indices

...
sorted = Lambda(top_k, arguments={'k': 10})(prediction)

Here's a runnable test:

import numpy as np
import tensorflow as tf

from keras import Sequential
from keras.layers import Lambda

def top_k(input, k):
  return tf.nn.top_k(input, k=k, sorted=True).indices

model = Sequential()
model.add(Lambda(top_k, input_shape=(10,), arguments={'k': 10}))

data = np.array([
  [0, 5, 2, 1, 3, 6, 1, 2, 7, 4],
  [2, 4, 3, 1, 2, 0, 1, 5, 2, 4],
  [8, 9, 1, 8, 3, 0, 1, 3, 2, 6],
])

print(model.predict(x=data))
# Prints:
# [[8 5 1 9 4 2 7 3 6 0]
#  [7 1 9 2 0 4 8 3 6 5]
#  [1 0 3 9 4 7 8 2 6 5]]

The argument k can be any value less than the dimension you're sorting, e.g., for k=4 the output will be:

[[8 5 1 9]
 [7 1 9 2]
 [1 0 3 9]]