python - 保存使用 BatchNorm 的 Tensorflow 模型

Question

我正在尝试使用 Tensorflow 从 GAN 保存生成器模型。我正在使用的模型有几个批量归一化层。当我保存权重时，只有运行全局变量初始值设定项才能成功恢复它们，但我不必这样做，因为所有变量都正在恢复。如果我在恢复之前运行全局变量初始值设定项，当我使用加载的权重运行推理并为批规范参数设置 is_training=False 时，模型的性能会非常差。但是，如果 is_training=True，则模型将按预期执行。这种行为应该完全相反。

为了节省权重，我这样做:

t_vars = tf.trainable_variables()
g_vars = [var for var in t_vars if 'g_' in var.name]
g_saver = tf.train.Saver(g_vars)
... train model ...
g_saver.save(sess, "weights/generator/gen.ckpt")

当我恢复权重时，我使用相同的模型定义并执行以下操作:

t_vars = tf.trainable_variables()
g_vars = [var for var in t_vars if 'g_' in var.name]

init = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init)

g_saver = tf.train.Saver(g_vars)
g_saver.restore(sess, "./weights/generator/gen.ckpt")

您需要执行特殊程序来计算批处理标准权重吗？我是否缺少一些变量集合？

编辑:

我使用以下方法定义批量归一化层:

conv1_norm = tf.contrib.layers.batch_norm(conv1, is_training=training

我发现将variables_collections=["g_batch_norm_non_trainable"]添加到batch_norm函数中，然后执行

g_vars = list(set([var for var in t_vars if 'g_' in var.name] + tf.get_collection("g_batch_norm_non_trainable")))

有效，但是对于一个简单的减肥指令来说，这似乎相当复杂。

Answer 1

当您使用 tf.contrib.layers.batch_norm 定义批量标准化时使用默认参数(如您的情况所示)，将创建三个变量:beta、moving_mean 和 moving_variance。第一个是唯一的可训练变量，其他两个包含在 tf.GraphKeys.GLOBAL_VARIABLES 集合中。

这就是为什么使用以下行中的可训练变量定义的 g_vars 不会同时在列表中获得 moving_mean 和 moving_variance :

g_vars = [var for var in t_vars if 'g_' in var.name]

由于您似乎只想保存生成器变量，因此我建议使用变量范围来定义生成器网络。

对随机张量进行上采样并使用批量归一化的示例:

import tensorflow as tf
import numpy as np

input_layer = tf.placeholder(tf.float32, (2, 7, 7, 64))  # (batch, height, width, in_channels)

with tf.variable_scope('generator', reuse=tf.AUTO_REUSE):
    # define your generator network here ...
    t_conv_layer = tf.layers.conv2d_transpose(input_layer,
                    filters=32, kernel_size=[3, 3], strides=(2, 2), padding='SAME', name='t_conv_layer')

    batch_norm = tf.contrib.layers.batch_norm(t_conv_layer, is_training=True, scope='my_batch_norm')
    print(batch_norm) # Tensor("generator/my_batch_norm/FusedBatchNorm:0", shape=(2, 14, 14, 32), dtype=float32)

您可以通过打印来检查 tf.trainable_variables() 和 tf.global_variables() 的变量列表。由于可训练变量位于全局变量列表中，如所述 here ，我们可以将 g_vars 定义为:

g_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='generator')

如果我们检查此列表，我们将拥有我们想要的批量标准化的所有变量:

for var in g_vars:
    print("variable_name: {:45}, nb_parameters: {}".format(var.name, np.prod(var.get_shape().as_list())))

产生输出:

variable_name: generator/t_conv_layer/kernel:0              , nb_parameters: 18432
variable_name: generator/t_conv_layer/bias:0                , nb_parameters: 32
variable_name: generator/my_batch_norm/beta:0               , nb_parameters: 32
variable_name: generator/my_batch_norm/moving_mean:0        , nb_parameters: 32
variable_name: generator/my_batch_norm/moving_variance:0    , nb_parameters: 32