Question

我正在尝试在keras中定义一个自定义图层，但是我找不到一种方法来将tf.cond与layers.Lambda函数一起变形

        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

由于true函数的主体很大，因此您可以创建一个自定义图层，如下所示：

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.]]

如何用keras.layers.Lambda包装tf.cond函数？

1 个答案: