Как обернуть функцию tf.cond в keras.layers.Lambda?

Question

Я пытаюсь определить пользовательский слой в кератах, но не могу найти способ деформировать tf.cond со слоями. Функция лямбда

        matches = tf.cond(
            tf.greater(N, 0),
            lambda: match_boxes(
                anchors, groundtruth_boxes,
                positives_threshold=positives_threshold,
                negatives_threshold=negatives_threshold,
                force_match_groundtruth=True
            ),
            lambda: only_background
        )

Answer 1

Поскольку тело вашей истинной функции очень велико, вы можете создать собственный слой, подобный этому:

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):

  def __init__(self, **kwargs):
    super(CustomLayer, self).__init__()
    self.pred = kwargs.get('pred', False)

  def call(self, inputs):
    def true_fn(x):
      return x + 1.

    return tf.cond(self.pred,
                   true_fn=lambda: true_fn(inputs),
                   false_fn=lambda: tf.identity(inputs))

Тестирование:

inputs = tf.placeholder(tf.float32, shape=(None, 1))
pred = tf.placeholder(tf.bool, shape=())

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Dense(1, kernel_initializer=tf.initializers.ones))
model.add(CustomLayer(pred=pred))

outputs = model(inputs)

with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  print(outputs.eval({inputs: [[1.]], pred: False})) # [[1.]]
  print(outputs.eval({inputs: [[1.]], pred: True})) # [[2.]]