Tensorflowのif文でusing a `tf.Tensor` as a Python `bool` is not allowed:

tf.kerasでcustom layerを書いていた時、ある確率でランダムに水増しを行う関数を書いていた。

model.fitを実行した際に表題のエラーが発生。

OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

AutoGraphではif文を解釈出来ないらしい。

AutoGraphの条件分岐に関するチュートリアルを見てみる。

tf.function で性能アップ | TensorFlow Core

if文を書く代わりに、tf.condで条件分岐をするらしい。

tf.cond | TensorFlow v2.16.1

Return true_fn() if the predicate pred is true else false_fn().

tf.cond(
pred, true_fn=None, false_fn=None, name=None
)
Args
pred A scalar determining whether to return the result of true_fn or false_fn.
true_fn The callable to be performed if pred is true.
false_fn The callable to be performed if pred is false.
name Optional name prefix for the returned tensors.

if文の代わりにtf.condを使ってcustom layerを書くと、上手くmodel.fit出来るようになった。

下記の例では75%の確率で入力を２倍するレイヤーとして働く。

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

  def __init__(self, **kwargs):
    super(CustomLayer, self).__init__()

  def call(self, inputs):
    def true_fn(x):
      return x * 2

    prob = tf.random.uniform([], 0, 1) < 0.75
    
    return tf.cond(prob,
                   true_fn=lambda: true_fn(inputs),
                   false_fn=lambda: tf.identity(inputs))

参考

How to wrap tf.cond function with keras.layers.Lambda?

I'm trying to define a custom layer in keras,but I can't find a way to warp tf.cond with layers.Lambda function matches ...