Batch Renormalization with Tensorflow 2+ and tf.function

Question

I'm trying to use Batch Renormalization in my model which is implemented in tf.keras.layers.BatchNormalization via the renorm_clipping parameter:

A dictionary that may map keys 'rmax', 'rmin', 'dmax' to scalar Tensors used to clip the renorm correction. The correction (r,d) is used as corrected_value = normalized_value * r + d, with r clipped to [rmin, rmax], and d to [-dmax, dmax]. Missing rmax, rmin, dmax are set to inf, 0, inf, respectively.

However, this minimal code example fails with Tensorflow 2.3.1:

# The actual input we care about
inputs = tf.keras.Input([1])

# The current global step
step = tf.keras.Input([1])

# Compute renorm parameters from global step
rmax = 1 + step / 1000
rmin = 1 / rmax
dmax = step / 1000

# Instantiate batch norm layer
layer = tf.keras.layers.BatchNormalization(renorm=True, renorm_clipping={
    "rmax": rmax,
    "rmin": rmin,
    "dmax": dmax,
})(inputs)

# Build Model
model = tf.keras.Model(inputs=[inputs, step], outputs=layer)

# Use model in tf.function
@tf.function
def predict(x):
    return model(x, training=True)

predict([tf.constant([[1.0], [1.0]]), tf.constant([[1.0], [2.0]])])

It fails in the predict() with:

_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [, , ]

If I set step to a constant value, everything works as expected, so clearly the problem is that I pass the Input tensor as a parameter to the BatchNormalization constructor. But the batch norm parameters need to updated dynamically and adding the step as an additional input seemed like the most elegant solution.

How would I get this example to work?

Batch Renormalization with Tensorflow 2+ and tf.function

0 Answers