TensorFlow中多神经网络共享权重矩阵与偏置的最优实现方法

嘿，这个场景我在做多任务学习的时候刚好实践过，TensorFlow里实现多网络共享权重和偏置有几种靠谱的路子，但最优的得看你用的是TF2.x还是TF1.x，下面给你拆解最实用的方案：

在TF2.x的动态图模式下，最直观且易维护的是两种方式：显式复用变量（适合简单场景）和封装共享层（适合复杂网络）。

方法1：显式复用变量（简单场景首选） #

直接定义好要共享的权重矩阵和偏置变量，然后在不同的网络结构里直接引用它们——完全透明，调试起来特别方便：

import tensorflow as tf

# 定义全局共享的权重与偏置
shared_weights = tf.Variable(tf.random.normal([784, 256]), name="shared_dense_weights")
shared_bias = tf.Variable(tf.zeros([256]), name="shared_dense_bias")

# 第一个网络：任务A分类
def task_a_network(inputs):
    # 复用共享变量做特征提取
    x = tf.matmul(inputs, shared_weights) + shared_bias
    x = tf.nn.relu(x)
    # 任务A独有的输出层
    x = tf.keras.layers.Dense(10, activation="softmax")(x)
    return x

# 第二个网络：任务B分类
def task_b_network(inputs):
    # 完全复用同一组共享变量
    x = tf.matmul(inputs, shared_weights) + shared_bias
    x = tf.nn.relu(x)
    # 任务B独有的输出层
    x = tf.keras.layers.Dense(5, activation="softmax")(x)
    return x

# 测试调用
batch_a = tf.random.normal([32, 784])
batch_b = tf.random.normal([32, 784])
pred_a = task_a_network(batch_a)
pred_b = task_b_network(batch_b)

方法2：封装共享层（复杂网络首选） #

如果共享的是一整套层（比如多卷积+全连接的特征提取模块），把共享逻辑封装成自定义Layer或Model，然后只实例化一次，多个网络调用同一个实例即可：

import tensorflow as tf

# 封装共享的特征提取模块
class SharedFeatureExtractor(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(256, activation="relu")
        self.dense2 = tf.keras.layers.Dense(128, activation="relu")
        self.dropout = tf.keras.layers.Dropout(0.3)
    
    def call(self, inputs, training=False):
        x = self.dense1(inputs)
        x = self.dense2(x)
        x = self.dropout(x, training=training)
        return x

# 关键：只创建一个共享层实例
shared_extractor = SharedFeatureExtractor()

# 构建任务A的完整模型
def build_task_a_model():
    inputs = tf.keras.Input(shape=(784,))
    # 复用共享特征提取层
    features = shared_extractor(inputs)
    outputs = tf.keras.layers.Dense(10, activation="softmax")(features)
    return tf.keras.Model(inputs, outputs)

# 构建任务B的完整模型
def build_task_b_model():
    inputs = tf.keras.Input(shape=(784,))
    # 复用同一个共享层实例
    features = shared_extractor(inputs)
    outputs = tf.keras.layers.Dense(5, activation="softmax")(features)
    return tf.keras.Model(inputs, outputs)

# 实例化两个模型
model_a = build_task_a_model()
model_b = build_task_b_model()

# 验证变量共享：两个模型的共享层是同一个对象
print(model_a.layers[1].dense1 is model_b.layers[1].dense1)  # 输出 True

这种方式代码模块化程度高，后续修改共享逻辑只需要改一处，所有网络都会同步更新，非常适合大型项目。

如果还在维护TF1.x的静态图代码，**tf.variable_scope配合reuse=tf.AUTO_REUSE**是经典且省心的实现方式：

import tensorflow as tf

def shared_feature_layers(inputs):
    # 用variable_scope限定共享变量的命名空间
    with tf.variable_scope("shared_layers", reuse=tf.AUTO_REUSE):
        weights = tf.get_variable(
            name="dense_weights",
            shape=[784, 256],
            initializer=tf.random_normal_initializer()
        )
        bias = tf.get_variable(
            name="dense_bias",
            shape=[256],
            initializer=tf.zeros_initializer()
        )
        x = tf.matmul(inputs, weights) + bias
        x = tf.nn.relu(x)
        return x

# 任务A网络
input_a = tf.placeholder(tf.float32, shape=[None, 784])
features_a = shared_feature_layers(input_a)
output_a = tf.layers.dense(features_a, 10, activation="softmax")

# 任务B网络
input_b = tf.placeholder(tf.float32, shape=[None, 784])
features_b = shared_feature_layers(input_b)
output_b = tf.layers.dense(features_b, 5, activation="softmax")

# 查看共享变量：两个网络的共享变量是同一个实例
print(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, "shared_layers/dense_weights"))

tf.AUTO_REUSE会自动判断变量是否已创建，避免重复定义，是TF1.x里最不容易出错的共享方式。

• 显式可控：不管是直接复用变量还是封装共享层，你都能明确知道哪些部分是共享的，不会出现隐式共享导致的调试噩梦。
• 资源高效：共享变量意味着内存/显存中只存一份权重，训练时也只需要更新一次共享参数，不会浪费计算资源。
• 易维护扩展：代码结构清晰，后续新增任务或修改共享逻辑只需要改动一处，所有关联网络都会同步更新。

内容的提问来源于stack exchange，提问作者K. Project