使用 Tensorflow 2+ 和 tf.function 进行批量重整化

Question

我正在尝试在我的模型中使用批量重整化，这是tf.keras.layers.BatchNormalization通过renorm_clipping参数实现的：

一个字典，可以将键 'rmax'、'rmin'、'dmax' 映射到用于剪辑 renorm 校正的标量张量。校正 (r,d) 用作 Corrected_value = normalized_value * r + d，其中 r 剪裁为 [rmin, rmax]，而 d 剪裁为 [-dmax, dmax]。缺失的rmax、rmin、dmax分别设置为inf、0、inf。

然而，这个最小的代码示例在 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]])])

它在以下情况下失败predict()：

_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'RealDiv_7:0' shape=(None, 1) dtype=float32>, <tf.Tensor 'AddV2_2:0' shape=(None, 1) dtype=float32>, <tf.Tensor 'RealDiv_8:0' shape=(None, 1) dtype=float32>]

如果我设置step为一个常量值，一切都会按预期工作，所以很明显问题是我将Input张量作为参数传递给BatchNormalization构造函数。但是批处理规范参数需要动态更新，并将其添加step为附加输入似乎是最优雅的解决方案。

我如何让这个例子工作？