Keras通过argmax输出单个值

Question

我正在努力在Keras建立一个非常简单的神经网络：

model = Sequential()
model.add(Dense(40, input_dim=186, activation='relu', name='x'))
model.add(Dense(3, activation='softmax'))

这有效，并输出三维矢量（例​​如0 1 0 ）。 我想添加一个使用argmax发送单个值的图层，而不是这个向量。

我认为这样可行：

model.add(Lambda(lambda x: K.cast(K.argmax(x), dtype='float32')))

但这次抛出（5962是训练样本的数量）：

ValueError: Error when checking target: expected lambda_1 to have 1 dimensions, but got array with shape (5962, 3)

我怎么做到这一点？

请注意，我在模型中将此视为实际的ArgMax层，类似于TensorFlow的ArgMax 。

Answer 1

感谢@today指出我正确的方向。 您应该在训练后添加图层并且一切正常：

model = Sequential()
model.add(Dense(40, input_dim=186, activation='relu', name='x'))
model.add(Dense(classes, activation='softmax'))

model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
model.fit(X_train, Y_train, batch_size=50, epochs=100, validation_data=(X_test, Y_test))
model.add(Lambda(lambda x: K.cast(K.argmax(x), dtype='float32'), name='y_pred'))
model.save('data/trained.h5')

现在，这将把ArgMax图层添加到模型中！