tensorflow - softmax（keras）之后的单个神经元层

Question

我需要创建一个神经网络（使用 keras），它的最后一层有一个神经元，其中包含神经元的索引，在前面的 softmax 层中具有最大值预测。

例如，我的 softmax 层给出了这样的结果：

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

我希望单个神经元层（在 softmax 层之后）给出结果2（考虑基于 0 的估值）。

我怎样才能做到这一点？

Answer 1

使用来自@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)))