Keras выводит одно значение через argmax

Question

Я пытаюсь построить действительно простую нейронную сеть в Керасе:

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, аналогично ArgMax TensorFlow .

Answer 1

Спасибо @ сегодня за указание в правильном направлении. Вы должны добавить слой после тренировки, и все в порядке:

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!