TensorFlow 2.0中如何获取预制Dense层变量用于原生训练流程？

嘿，这个问题我之前也踩过坑！TF2放弃了TF1里全局的tf.trainable_variables()，但其实针对使用tf.keras.layers.Dense这类层的场景，有几种很顺手的方式来获取可训练变量，完美适配你用原生TF函数写训练流程的需求：

方法1：直接保存层实例，逐个获取变量 #

如果你只是用了少数几个独立的Keras层，最简单的方式就是在创建层的时候把它们存成变量，之后直接访问每个层的trainable_variables属性，把所有变量合并起来：

# 先定义并保存你的Dense层实例
dense_hidden = tf.keras.layers.Dense(units=64, activation='relu')
dense_output = tf.keras.layers.Dense(units=1, activation='sigmoid')

def logistic_regression(x):
    x = dense_hidden(x)
    return dense_output(x)

# 优化流程里收集所有可训练变量
def run_optimization(x, y):
    with tf.GradientTape() as g:
        pred = logistic_regression(x)
        loss = cross_entropy(pred, y)
    # 合并所有层的可训练变量
    all_trainable_vars = dense_hidden.trainable_variables + dense_output.trainable_variables
    # 计算梯度并更新
    gradients = g.gradient(loss, all_trainable_vars)
    optimizer.apply_gradients(zip(gradients, all_trainable_vars))

每个Keras层的trainable_variables会返回一个列表，包含它的权重（kernel）和偏置（bias），把多个层的列表相加就能得到所有需要优化的变量。

方法2：用tf.Module或keras.Model封装层（推荐） #

如果你的模型有多个层，逐个合并变量太麻烦，推荐把所有层封装到tf.Module或者tf.keras.Model里——这两个类都是TF2里用来组织模型结构和变量的核心工具，它们会自动收集所有子模块（也就是你的Dense层）的可训练变量：

用tf.Module的示例： #

class LogisticRegressionModel(tf.Module):
    def __init__(self):
        super().__init__()
        self.dense_hidden = tf.keras.layers.Dense(units=64, activation='relu')
        self.dense_output = tf.keras.layers.Dense(units=1, activation='sigmoid')
    
    def __call__(self, x):
        x = self.dense_hidden(x)
        return self.dense_output(x)

# 实例化模型
model = LogisticRegressionModel()

def run_optimization(x, y):
    with tf.GradientTape() as g:
        pred = model(x)
        loss = cross_entropy(pred, y)
    # 直接获取整个模型的所有可训练变量
    gradients = g.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

用keras.Model的示例（其实和Module逻辑一致，只是多了一些Keras的额外功能）： #

class LogisticRegressionModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense_hidden = tf.keras.layers.Dense(units=64, activation='relu')
        self.dense_output = tf.keras.layers.Dense(units=1, activation='sigmoid')
    
    def call(self, x):
        x = self.dense_hidden(x)
        return self.dense_output(x)

model = LogisticRegressionModel()

# 优化函数和上面完全一样
def run_optimization(x, y):
    with tf.GradientTape() as g:
        pred = model(x)
        loss = cross_entropy(pred, y)
    gradients = g.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

这种方式的好处是不管你加多少层，都只需要访问model.trainable_variables就能拿到所有需要优化的变量，非常整洁。

避坑提醒：别用兼容模式 #

虽然TF2提供了tf.compat.v1.trainable_variables()来兼容TF1的写法，但不推荐长期使用——这个函数会返回全局范围内的所有可训练变量，当你的代码里有多个模型或者变量集合时，很容易出现混淆，而且后续版本可能会逐步移除兼容接口。

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