我想對批處理規范化層的變量添加條件操作。具體來說，在浮點數中訓練，然后在微調二次訓練階段量化。為此，我想在變量上添加一個 tf.cond 操作（均值和變量的比例、移位和 exp 移動平均值）。我用tf.layers.batch_normalization 我寫的batchnorm層替換了它（見下文）。這個函數工作得很好（即我用兩個函數得到了相同的度量），我可以向變量添加任何管道（在batchnorm 操作之前）。問題是性能（運行時）急劇下降（即通過簡單地用我自己的函數替換 layer.batchnorm 有一個 x2 因子，如下所述）。def batchnorm(self, x, name, epsilon=0.001, decay=0.99):    epsilon = tf.to_float(epsilon)    decay = tf.to_float(decay)    with tf.variable_scope(name):        shape = x.get_shape().as_list()        channels_num = shape[3]        # scale factor        gamma = tf.get_variable("gamma", shape=[channels_num], initializer=tf.constant_initializer(1.0), trainable=True)        # shift value        beta = tf.get_variable("beta", shape=[channels_num], initializer=tf.constant_initializer(0.0), trainable=True)        moving_mean = tf.get_variable("moving_mean", channels_num, initializer=tf.constant_initializer(0.0), trainable=False)        moving_var = tf.get_variable("moving_var", channels_num, initializer=tf.constant_initializer(1.0), trainable=False)        batch_mean, batch_var = tf.nn.moments(x, axes=[0, 1, 2]) # per channel        update_mean = moving_mean.assign((decay * moving_mean) + ((1. - decay) * batch_mean))        update_var = moving_var.assign((decay * moving_var) + ((1. - decay) * batch_var))        tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_mean)        tf.add_to_collection(tf.GraphKeys.UPDATE_OPS, update_var)        bn_mean = tf.cond(self.is_training, lambda: tf.identity(batch_mean), lambda: tf.identity(moving_mean))        bn_var = tf.cond(self.is_training, lambda: tf.identity(batch_var), lambda: tf.identity(moving_var))        with tf.variable_scope(name + "_batchnorm_op"):            inv = tf.math.rsqrt(bn_var + epsilon)            inv *= gamma            output = ((x*inv) - (bn_mean*inv)) + beta    return output我將不勝感激以下任何問題的幫助：關于如何提高我的解決方案的性能（減少運行時間）的任何想法？是否可以在 batchnorm 操作之前將我自己的運算符添加到 layer.batchnorm 的變量管道中？有沒有其他解決方案可以解決同樣的問題？