TensorBoard中MNIST示例生成意外Conv2d_1与Dense_1层的原因

这个问题我之前也碰到过，核心原因是tf.layers的API在Estimator的model_fn中被多次调用时，自动生成了新的层实例，具体来说有两个关键点：

1. 自动命名机制导致重复创建层 #

当你使用tf.layers.conv2d或者tf.layers.dense这类高阶API时，如果没有显式指定name参数，TensorFlow会自动给每个层分配递增的名称（比如第一次调用生成conv2d，第二次就变成conv2d_1）。而在Estimator框架中，model_fn会被多次调用——比如训练阶段、评估阶段（甚至预测阶段），每次调用都会重新执行你的层定义代码，这就会创建出名称带后缀的新层，反映在TensorBoard里就是多出了conv2d_1和dense_1。

2. 变量复用未开启 #

Estimator在不同模式下（TRAIN/EVAL/PREDICT）会共享同一个计算图，但如果你的层没有设置变量复用，每次调用model_fn时都会重新创建新的变量和对应的层节点，而非复用已有的。

解决方法 #

方案一：给每个层指定固定的name参数，并开启变量复用 #

修改你的层定义代码，给每个层加上明确的name，同时在model_fn中通过变量作用域来控制复用：

def model_fn(features, labels, mode):
    input_layer = tf.reshape(features["x"], [-1, 28, 28, 1])
    
    # 使用变量作用域，设置reuse参数
    with tf.variable_scope('cnn_layers', reuse=mode != tf.estimator.ModeKeys.TRAIN):
        conv1 = tf.layers.conv2d(
            inputs=input_layer,
            filters=32,
            kernel_size=[5, 5],
            padding="same",
            activation=tf.nn.relu,
            name='conv1'  # 指定固定名称
        )
        pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2, name='pool1')
        conv2 = tf.layers.conv2d(
            inputs=pool1,
            filters=64,
            kernel_size=[5, 5],
            padding="same",
            activation=tf.nn.relu,
            name='conv2'  # 指定固定名称
        )
        pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2, name='pool2')
    
    with tf.variable_scope('fc_layers', reuse=mode != tf.estimator.ModeKeys.TRAIN):
        pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
        dense = tf.layers.dense(
            inputs=pool2_flat,
            units=1024,
            activation=tf.nn.relu,
            name='dense'  # 指定固定名称
        )
        dropout = tf.layers.dropout(
            inputs=dense,
            rate=0.4,
            training=mode == tf.estimator.ModeKeys.TRAIN,
            name='dropout'
        )
        logits = tf.layers.dense(
            inputs=dropout,
            units=10,
            name='logits'  # 指定固定名称
        )
    
    # 后续的损失、优化器等定义...

方案二：使用tf.make_template包装层定义 #

这个方法更简洁，它会自动帮你处理变量复用，确保每次调用时都使用同一个层实例：

def build_model(inputs, is_training):
    conv1 = tf.layers.conv2d(inputs=inputs, filters=32, kernel_size=[5,5], padding="same", activation=tf.nn.relu)
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2,2], strides=2)
    conv2 = tf.layers.conv2d(inputs=pool1, filters=64, kernel_size=[5,5], padding="same", activation=tf.nn.relu)
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2,2], strides=2)
    pool2_flat = tf.reshape(pool2, [-1, 7*7*64])
    dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
    dropout = tf.layers.dropout(inputs=dense, rate=0.4, training=is_training)
    logits = tf.layers.dense(inputs=dropout, units=10)
    return logits

# 用make_template包装模型构建函数
build_model_template = tf.make_template('mnist_model', build_model)

def model_fn(features, labels, mode):
    input_layer = tf.reshape(features["x"], [-1, 28, 28, 1])
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)
    logits = build_model_template(input_layer, is_training)
    
    # 后续的损失、优化器等定义...

这样不管model_fn被调用多少次，都会复用同一个模板里的层，不会生成带后缀的新层，TensorBoard里就只会显示你定义的2个卷积和2个全连接层了。

内容的提问来源于stack exchange，提问作者air-fox