Tensorflow softmax后的单神经元层（keras）

我需要创建一个神经网络（带有keras），该网络的最后一层是一个神经元，其中包含神经元的索引，该神经元的最大值预测位于前一个softmax层

例如，我的softmax图层的结果如下：

[0.1, 0.1, 0.7, 0.0, 0.05, 0.05]

我希望单神经元层（在softmax层之后）给出结果

2

（考虑基于0的值）

---

我该怎么做？

使用@Eric Platon的想法：

import keras.backend as K
K.argmax(x, axis=-1)

但我不确定是否可以使用后端函数作为层。可能需要将其包裹在lambda层中：

from keras.layers import Lambda
model.add(Lambda(lambda x: K.argmax(x, axis=-1)))

在TF中，它是Argmax:。也许是Keras的一个好关键词。