Есть ли внутри модели Keras способ скрыть некоторые недопустимые результаты и перенормировать логиты?

Question

Я использую функциональный API Keras. У меня есть некоторая модель, которая выводит распределение вероятностей с помощью слоя softmax:

action_logits = Dense(units=self.action_space, activation='softmax')(prev_layer)

Затем я маскирую недопустимые действия (или классы, если хотите), умножая логиты на битовый вектор, представляющий законный actions:

mask_illegal_moves = keras.layers.multiply([action_logits, valid_actions])

Наконец, я хочу перенормировать логиты, теперь, когда я установил вывод для некоторых действий равным 0. Это кажется очень простым делом, но я не могу получить это на работу. Например, другой слой softmax не дал желаемых результатов. Более того, поиск в Google любого слоя «нормализации» в основном привел меня к BatchNorm, который меня здесь не интересует.

Любые советы будут с благодарностью приняты!

Answer 1

Я должен был уточнить мой первоначальный вопрос. Смотрите полезный ответ от thushv89. Я нашел решение, которое заключалось в использовании лямбда-слоя:

action_probs = Dense(units=self.action_space, activation='softmax')(skip_2)
action_probs_masked = Multiply()([action_probs, valid_actions])
layer = Lambda(lambda x: x / keras.backend.sum(x, axis=1)[:,None])
actions = layer(action_probs_masked)

Answer 2

Вы можете сделать следующее,

action_logits = Dense(units=self.action_space)(prev_layer)
action_logits_masked = Multiply()([action_logits, valid_actions])
action_probs = Activation('softmax')(action_logits_masked)

Объяснено:

• Сначала мы получаем логиты (обратите внимание, что логиты - это то, что они называются до применяя софтмакс). Так что не используйте activation='softmax'
• Примените маску, используя Multiply layer
• Примените softmax к маскированному слою.