python - 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's 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 层添加到模型中!