如何使用tf.layers.conv2d实现权重共享并分阶段训练自编码器

Hey，我来帮你搞定这个用tf.layers.conv2d实现分阶段训练自编码器的问题！其实TensorFlow的高层conv2d层完全能支持你要的分阶段训练、输入切换和权重共享，咱们一步步来拆解：

一、分阶段训练：通过trainable参数+变量作用域控制 #

tf.layers.conv2d虽然封装了变量创建，但可以通过trainable参数精准控制哪些层参与梯度更新，配合变量作用域的复用机制，就能实现“外层→中层→内层”的分阶段训练。

首先，咱们把自编码器拆成清晰的层级结构（编码器：外层conv1→中层conv2→内层conv3；解码器：内层conv4→中层conv5→外层conv6），然后写一个可配置的构建函数：

import tensorflow as tf

def build_autoencoder(inputs, trainable_layers=None):
    # 默认所有层都可训练，传入指定列表则只让这些层参与训练
    if trainable_layers is None:
        trainable_layers = ['conv1', 'conv2', 'conv3', 'conv4', 'conv5', 'conv6']
    
    # 编码器
    with tf.variable_scope('encoder', reuse=tf.AUTO_REUSE):
        # 外层conv1
        conv1 = tf.layers.conv2d(
            inputs, 32, 3, padding='same', activation=tf.nn.relu,
            trainable='conv1' in trainable_layers, name='conv1'
        )
        # 中层conv2
        conv2 = tf.layers.conv2d(
            conv1, 16, 3, padding='same', activation=tf.nn.relu,
            trainable='conv2' in trainable_layers, name='conv2'
        )
        # 内层conv3（瓶颈层）
        conv3 = tf.layers.conv2d(
            conv2, 8, 3, padding='same', activation=tf.nn.relu,
            trainable='conv3' in trainable_layers, name='conv3'
        )
    
    # 解码器
    with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
        # 内层conv4
        conv4 = tf.layers.conv2d_transpose(
            conv3, 16, 3, padding='same', activation=tf.nn.relu,
            trainable='conv4' in trainable_layers, name='conv4'
        )
        # 中层conv5
        conv5 = tf.layers.conv2d_transpose(
            conv4, 32, 3, padding='same', activation=tf.nn.relu,
            trainable='conv5' in trainable_layers, name='conv5'
        )
        # 外层conv6（输出层）
        conv6 = tf.layers.conv2d_transpose(
            conv5, 3, 3, padding='same', activation=tf.nn.sigmoid,
            trainable='conv6' in trainable_layers, name='conv6'
        )
    
    return conv6

接下来分阶段训练就很简单了：

1. 第一阶段：训练外层
   只让conv1和conv6参与训练，其他层固定：

inputs = tf.placeholder(tf.float32, [None, 64, 64, 3])
# 指定仅外层可训练
outputs = build_autoencoder(inputs, trainable_layers=['conv1', 'conv6'])
loss = tf.losses.mean_squared_error(inputs, outputs)
optimizer = tf.train.AdamOptimizer(1e-3).minimize(loss)
# 后续就是常规的训练循环，这里省略

2. 第二阶段：训练中层
   固定外层，让conv2和conv5参与训练：

outputs = build_autoencoder(inputs, trainable_layers=['conv2', 'conv5'])
loss = tf.losses.mean_squared_error(inputs, outputs)
optimizer = tf.train.AdamOptimizer(1e-3).minimize(loss)

3. 第三阶段：训练内层
   固定外层和中层，让conv3和conv4参与训练：

outputs = build_autoencoder(inputs, trainable_layers=['conv3', 'conv4'])
loss = tf.losses.mean_squared_error(inputs, outputs)
optimizer = tf.train.AdamOptimizer(1e-3).minimize(loss)

二、切换变量作用域调整输入（跳过内层） #

如果想在训练外层时直接跳过中层和内层，让编码器的输出直接对接解码器的外层，可以通过阶段分支+变量作用域复用实现。修改构建函数，加入stage参数控制网络连接：

def build_autoencoder_stage(inputs, stage):
    with tf.variable_scope('autoencoder', reuse=tf.AUTO_REUSE):
        if stage == 'outer':
            # 跳过中层、内层，直接外层编码器→外层解码器
            conv1 = tf.layers.conv2d(inputs, 32, 3, padding='same', activation=tf.nn.relu, name='conv1')
            conv6 = tf.layers.conv2d_transpose(conv1, 3, 3, padding='same', activation=tf.nn.sigmoid, name='conv6')
            return conv6
        elif stage == 'middle':
            # 固定外层，训练中层：编码器conv1→conv2，解码器conv5→conv6
            conv1 = tf.layers.conv2d(inputs, 32, 3, padding='same', activation=tf.nn.relu, trainable=False, name='conv1')
            conv2 = tf.layers.conv2d(conv1, 16, 3, padding='same', activation=tf.nn.relu, name='conv2')
            # 跳过内层，conv2直接对接conv5
            conv5 = tf.layers.conv2d_transpose(conv2, 32, 3, padding='same', activation=tf.nn.relu, name='conv5')
            conv6 = tf.layers.conv2d_transpose(conv5, 3, 3, padding='same', activation=tf.nn.sigmoid, trainable=False, name='conv6')
            return conv6
        elif stage == 'inner':
            # 完整结构，固定外层、中层，训练内层
            conv1 = tf.layers.conv2d(inputs, 32, 3, padding='same', activation=tf.nn.relu, trainable=False, name='conv1')
            conv2 = tf.layers.conv2d(conv1, 16, 3, padding='same', activation=tf.nn.relu, trainable=False, name='conv2')
            conv3 = tf.layers.conv2d(conv2, 8, 3, padding='same', activation=tf.nn.relu, name='conv3')
            conv4 = tf.layers.conv2d_transpose(conv3, 16, 3, padding='same', activation=tf.nn.relu, name='conv4')
            conv5 = tf.layers.conv2d_transpose(conv4, 32, 3, padding='same', activation=tf.nn.relu, trainable=False, name='conv5')
            conv6 = tf.layers.conv2d_transpose(conv5, 3, 3, padding='same', activation=tf.nn.sigmoid, trainable=False, name='conv6')
            return conv6

调用时只需指定阶段，比如训练外层就用build_autoencoder_stage(inputs, 'outer')，这样就自动跳过了中间层。

三、tf.layers.conv2d实现权重共享 #

权重共享的核心是让不同层复用同一个变量集合，tf.layers.conv2d可以通过变量作用域的reuse参数或者**tf.AUTO_REUSE**实现：

场景1：不同输入共享同一卷积层 #

比如两个输入要经过相同的卷积处理：

def shared_conv_layer(inputs):
    with tf.variable_scope('shared_conv', reuse=tf.AUTO_REUSE):
        return tf.layers.conv2d(inputs, 32, 3, padding='same', activation=tf.nn.relu, name='conv')

# 两个不同输入复用同一卷积层的权重
output1 = shared_conv_layer(input1)
output2 = shared_conv_layer(input2)

这里tf.AUTO_REUSE会自动判断变量是否已存在，存在则复用，不存在则创建，非常方便。

场景2：自编码器对称层共享权重 #

如果想让编码器的conv1和解码器的conv5（对称层）共享权重，需要注意普通卷积和转置卷积的权重形状差异：普通卷积权重是[h, w, in_ch, out_ch]，转置卷积是[h, w, out_ch, in_ch]，所以要转置权重的通道维度：

with tf.variable_scope('encoder', reuse=tf.AUTO_REUSE):
    conv1 = tf.layers.conv2d(inputs, 32, 3, padding='same', name='conv1')
    # 获取conv1的权重和偏置
    conv1_kernel = tf.get_variable('conv1/kernel')
    conv1_bias = tf.get_variable('conv1/bias')

with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
    # 转置卷积使用conv1权重的转置版本
    conv5 = tf.layers.conv2d_transpose(
        decoder_input, 3, 3, padding='same', activation=tf.nn.relu,
        # 转置权重的通道维度
        kernel_initializer=tf.constant_initializer(tf.transpose(conv1_kernel, [0,1,3,2])),
        bias_initializer=tf.constant_initializer(conv1_bias),
        trainable=False, name='conv5'
    )

如果想让这两个层的权重同步更新，可以把conv5的trainable设为False，只优化conv1的权重，conv5会自动复用更新后的权重。

内容的提问来源于stack exchange，提问作者user3687879