我想在批量标准化层的变量上添加条件操作。具体来说,进行浮点训练,然后在微调二次训练阶段进行量化。为此,我想在变量上添加 tf.cond 操作(均值和 var 的缩放、移位和 exp 移动平均值)。
我用我编写的batchnorm层替换了tf.layers.batch_normalization
(见下文)。
这个函数工作完美(即我使用两个函数得到相同的指标),并且我可以将任何管道添加到变量中(在批标准化操作之前)。 问题是性能(运行时)急剧下降(即,通过简单地将layers.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操作之前将我自己的运算符添加到layers.batchnorm的变量管道中?
- 对于同一问题还有其他解决方案吗?
最佳答案
tf.nn.fused_batch_norm
经过优化并达到了目的。
我必须创建两个子图,每个模式一个,因为 fused_batch_norm
的接口(interface)不采用条件训练/测试模式(is_training 是 bool 而不是张量,因此它的图不是条件的) )。我在之后添加了条件(见下文)。然而,即使有两个子图,它的运行时间也与 tf.layers.batch_normalization 大致相同。
这是最终的解决方案(我仍然感谢任何改进意见或建议):
def batchnorm(self, x, name, epsilon=0.001, decay=0.99):
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)
(output_train, batch_mean, batch_var) = tf.nn.fused_batch_norm(x,
gamma,
beta, # pylint: disable=invalid-name
mean=None,
variance=None,
epsilon=epsilon,
data_format="NHWC",
is_training=True,
name="_batchnorm_op")
(output_test, _, _) = tf.nn.fused_batch_norm(x,
gamma,
beta, # pylint: disable=invalid-name
mean=moving_mean,
variance=moving_var,
epsilon=epsilon,
data_format="NHWC",
is_training=False,
name="_batchnorm_op")
output = tf.cond(self.is_training, lambda: tf.identity(output_train), lambda: tf.identity(output_test))
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)
return output
关于python - tensorflow 中的批量标准化: variables and performance,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/55474712/