如何在 Tensorflow 中处理像 tf.nn.fused_batch_norm 这样的 bool 值需要的操作

How to deal with bool valued needed operation like tf.nn.fused_batch_norm in Tensorflow

对于像tf.layers.batch_normalization这样的函数,很容易使用占位符train_flag作为training参数的输入,在我们定义了整个网络之后,我们可以输入TrueFalse for train_flag 在训练或推理阶段。

但是,对于像 tf.nn.fused_batch_norm(在 nn_impl.py 中定义)这样的操作,它只接受 python bool 作为参数 is_training 的输入,这是否意味着我们将需要在训练和推理阶段用不同 is_training 个参数构建网络两次?

tf.keras.layers.BatchNormalization

的定义文件中
def _fused_batch_norm(self, inputs, training):
    """Returns the output of fused batch norm."""

    def _fused_batch_norm_training():
        return nn.fused_batch_norm(
            inputs,
            self.gamma,
            self.beta,
            epsilon=self.epsilon)

    def _fused_batch_norm_inference():
        return nn.fused_batch_norm(
            inputs,
            self.gamma,
            self.beta,
            mean=self.moving_mean,
            variance=self.moving_variance,
            epsilon=self.epsilon,
            is_training=False)

    output, mean, variance = tf_utils.smart_cond(
        training, _fused_batch_norm_training, _fused_batch_norm_inference)
    if not self._bessels_correction_test_only:
        # Remove Bessel's correction to be consistent with non-fused batch norm.
        # Note that the variance computed by fused batch norm is
        # with Bessel's correction.
        sample_size = math_ops.cast(
                array_ops.size(inputs) / array_ops.size(variance), variance.dtype)
        factor = (sample_size - math_ops.cast(1.0, variance.dtype)) / sample_size
        variance *= factor

    training_value = tf_utils.constant_value(training)
    if training_value is None:
        momentum = tf_utils.smart_cond(training,
                                       lambda: self.momentum,
                                       lambda: 1.0)
    else:
        momentum = ops.convert_to_tensor(self.momentum)
    if training_value or training_value is None:
        mean_update = self._assign_moving_average(self.moving_mean, mean,
                                                  momentum)
        variance_update = self._assign_moving_average(self.moving_variance,
                                                      variance, momentum)
        self.add_update(mean_update, inputs=True)
        self.add_update(variance_update, inputs=True)

    return output

它利用training_value = tf_utils.constant_value(training)然后使用if training_value is None,其中tf_utils.constant_value应该是tf.contrib.util.constant_value