Keras output single value through argmax

Question

I'm trying to build a really simple neural network in Keras:

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

This works, and outputs a three-dimensional vector (eg 0 1 0 ). I'd like to add a layer that uses argmax to send out a single value, rather than this vector.

I figured this would work:

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

But this throws (5962 is the number of training samples):

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

How would I achieve this?

Note that I'd like this in the model as an actual ArgMax layer, similar to TensorFlow's ArgMax .

Answer 1

Thanks to @today for pointing me in the right direction. You should add the layer after training and all is fine:

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')

This will now have added the ArgMax layer to the model!