TensorFlow 2.0训练模型权重提取及模型保存加载问题咨询

我来帮你解决这两个核心问题——怎么选模型保存方式，以及如何正确提取权重和偏置。

一、save_weights vs save：该选哪一个？ #

先明确两者的核心区别，再结合你的需求给出建议：

• save_weights：只保存模型的权重参数（包括权重和偏置），不保存模型结构。
  • 优点：文件体积小，加载时需要先实例化你的MyModel类，适合你这种已经有固定模型结构、需要后续重新训练的场景。完全满足你“重新训练+提取权重分析”的需求。
• save：保存完整的模型（结构+权重+优化器状态+训练配置等），后续可以直接用tf.keras.models.load_model加载，不需要重新定义MyModel类。
  • 适用场景：如果你需要在没有模型类代码的环境中部署或加载模型，或者想一次性保存所有训练相关状态。但对于你的需求来说，save_weights足够灵活且轻便。

结论：既然你已经有明确的模型类定义，且需要重新训练+提取权重，用save_weights就够了。如果想做冗余备份，也可以同时用save保存完整模型，但日常训练用save_weights更高效。

二、如何正确提取模型的权重和偏置？ #

你用tf.compat.v1.get_collection没效果，是因为这是TensorFlow 1.x图模式的用法，而TF2默认是即时执行（Eager Execution）模式，直接用模型的内置属性就能轻松获取权重。

第一步：正确加载模型权重 #

首先要注意，TF2采用延迟构建模型的机制，实例化MyModel后，需要先构建模型（指定输入形状），再加载权重，否则模型参数还未初始化，加载会出问题：

from tensorflow.keras import Model, layers

# 重新定义你的模型类（和训练时一致）
class MyModel(Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = layers.Conv2D(filters=32, kernel_size=3, strides=[2,2], activation='relu')
        self.flatten = layers.Flatten()
        self.d1 = layers.Dense(units=64, activation="relu")
        self.d2 = layers.Dense(units=10, activation="softmax")
    def call(self, x):
        x = self.conv1(x)
        x = self.flatten(x)
        x = self.d1(x)
        x = self.d2(x)
        return x

# 实例化模型并构建输入形状（替换成你实际的输入尺寸，比如MNIST是(28,28,1)）
model = MyModel()
model.build(input_shape=(None, 28, 28, 1))  # None表示批量大小可变

# 加载训练好的权重
checkpoint_path = "./logs/model.ckpt"
model.load_weights(checkpoint_path)

第二步：提取权重和偏置 #

有两种常用方式，按需选择：

方式1：获取所有可训练变量（权重+偏置）

直接通过model.trainable_variables获取所有可训练参数，遍历即可查看或保存：

trainable_vars = model.trainable_variables
for var in trainable_vars:
    print(f"变量名称: {var.name}, 形状: {var.shape}")
    # 将变量转换为numpy数组，用于后续分析
    var_value = var.numpy()
    # 示例：保存到本地文件（可选）
    # np.save(f"./weights/{var.name.replace('/', '_')}.npy", var_value)

方式2：按层精准提取

如果你需要指定某一层的权重或偏置，直接通过模型的层属性获取：

# 获取Conv2D层的权重和偏置
conv1_weights = model.conv1.kernel.numpy()  # 卷积核权重
conv1_bias = model.conv1.bias.numpy()        # 卷积层偏置

# 获取全连接层d1的权重和偏置
d1_weights = model.d1.kernel.numpy()         # Dense层的权重矩阵
d1_bias = model.d1.bias.numpy()              # Dense层的偏置向量

# 获取d2层的参数同理
d2_weights = model.d2.kernel.numpy()
d2_bias = model.d2.bias.numpy()

这样就能轻松拿到所有你需要的权重和偏置，用于后续分析或可视化啦。

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