Keras中加载CNN前几层权重及安全删除预训练模型的方法

Great question! Let's break this down into actionable solutions based on your needs—whether you want to avoid loading the entire model, or need to safely clean up after loading it.

1. 无需加载整个模型：直接从HDF5文件提取指定层权重 #

Yes, you absolutely can extract weights for just the first 4 convolutional layers without loading the entire model. Keras .h5 files are actually HDF5 format archives that store layer weights in a structured way. You can use the h5py library to directly read only the weights you need, which is perfect for saving memory with large models.

Here's a step-by-step code example:

import h5py
import tensorflow as tf

# 1. 定义新模型，确保前4个卷积层与预训练模型结构完全匹配
# 包括滤波器数量、核大小、是否使用偏置等参数必须一致
new_model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu', name='conv2d_1'),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu', name='conv2d_2'),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu', name='conv2d_3'),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu', name='conv2d_4'),
    # 添加新任务所需的自定义层
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 2. 直接从.h5文件提取目标层权重
with h5py.File('keras_pretrained_model.h5', 'r') as h5_file:
    # 获取预训练模型的所有层名称
    layer_names = [name.decode('utf-8') for name in h5_file['layer_names']]
    
    # 筛选出前4个卷积层（可根据实际层名调整过滤逻辑）
    target_conv_layers = [name for name in layer_names if 'conv' in name][:4]
    
    # 将权重赋值给新模型的对应层
    for layer_name in target_conv_layers:
        # 获取新模型中匹配的层（可按名称或索引匹配）
        new_layer = new_model.get_layer(name=layer_name)
        
        # 读取预训练层的核权重和偏置
        kernel_weights = h5_file[f'/{layer_name}/kernel:0'][()]
        bias_weights = h5_file[f'/{layer_name}/bias:0'][()]
        
        # 为新层设置权重
        new_layer.set_weights([kernel_weights, bias_weights])

该方法的关键注意事项： #

• 层结构必须完全匹配：新模型的前4个卷积层，在滤波器数量、核大小、填充方式、是否启用偏置等参数上，必须和预训练模型的对应层完全一致，否则set_weights()会抛出错误。
• 层名称匹配：如果预训练模型使用自动生成的名称（如conv2d_1），确保新模型使用相同名称，或者改为按索引匹配层（比如new_model.layers[i]）。

2. 若必须加载整个模型：安全删除预训练模型且不影响新模型 #

如果你不确定预训练模型的精确层结构，更愿意先加载整个模型，也可以安全清理预训练模型而不丢失新模型。核心是先提取权重，再显式删除预训练模型并触发垃圾回收。

代码示例如下：

import tensorflow as tf
import gc

# 1. 加载完整的预训练模型
pretrained_model = tf.keras.models.load_model('keras_pretrained_model.h5')

# 2. 提取前4个卷积层的权重
# 根据你的模型层命名规则调整过滤逻辑
pretrained_conv_weights = [
    layer.get_weights() 
    for layer in pretrained_model.layers 
    if 'conv' in layer.name
][:4]

# 3. 定义新模型，确保前4个卷积层结构匹配
new_model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
    tf.keras.layers.Conv2D(128, (3, 3), activation='relu'),
    # 添加新任务层
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 4. 将预训练权重赋值给新模型
for i in range(4):
    new_model.layers[i].set_weights(pretrained_conv_weights[i])

# 5. 安全释放预训练模型占用的内存
# 先删除预训练模型对象
del pretrained_model

# 强制触发垃圾回收，立即释放内存
gc.collect()

# 注意：不要在这里调用tf.keras.backend.clear_session()！
# 它会清除整个Keras计算图，包括你的新模型。

重要提醒： #

如果确实需要调用tf.keras.backend.clear_session()（比如重置Keras状态），请先保存新模型到文件，之后再重新加载：

# 先保存新模型
new_model.save('temp_new_model.h5')
del new_model

# 清除会话
tf.keras.backend.clear_session()

# 重新加载新模型
new_model = tf.keras.models.load_model('temp_new_model.h5')

最终建议 #

• 优先使用HDF5直接读取法：处理大型模型时，这种方法内存占用最低，因为你从未将整个模型加载到内存中。
• 务必检查层兼容性：赋值权重前，哪怕是微小的结构不匹配（比如忘记启用偏置）都会导致错误。
• 避免不必要的clear_session调用：除非你已经保存了新模型，否则不要轻易使用这个方法。

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