如何保存Keras CNN模型用于迁移学习复用及加载后新增网络层

一、适配迁移学习的模型保存方式 #

迁移学习场景下推荐两种保存方案，可根据需求选择：

• 方案1：保存完整模型（结构+权重+编译配置）
  该方案无需保留原模型代码，加载后可直接修改适配新任务，是迁移学习的首选方案。训练完成后执行以下代码即可：

  # 保存为H5格式单文件
  model_cnn.save('cnn_pretrained.h5')
  # 也可保存为TensorFlow官方推荐的SavedModel文件夹格式，兼容性更好
  # model_cnn.save('cnn_pretrained_savedmodel')
  

• 方案2：仅保存权重
  适合需要灵活修改原模型底层结构的场景，需要自行保留原模型的结构代码，保存代码如下：

  model_cnn.save_weights('cnn_pretrained_weights.h5')
  

二、加载预训练模型并新增层实现迁移学习 #

2.1 加载完整预训练模型 #

如果采用方案1保存的完整模型，注意你原代码用到了Precision、Recall指标，需要通过custom_objects参数声明避免报错：

from tensorflow.keras.models import load_model
from tensorflow.keras.metrics import Precision, Recall

# 加载预训练模型
pretrained_model = load_model(
    'cnn_pretrained.h5',
    custom_objects={'Precision': Precision, 'Recall': Recall}
)

# 可选：冻结预训练模型的全部层，避免新任务训练时破坏已学习的特征
pretrained_model.trainable = False
# 如果仅需要冻结部分层，可指定层的trainable属性，比如只冻结前3层
# for layer in pretrained_model.layers[:3]:
#     layer.trainable = False

2.2 截取预训练模型的特征提取层 #

原模型最后两层是适配原二分类任务的全连接层，迁移学习时需要去掉这部分分类头，截取到你需要的特征输出层即可，比如要用到Flatten层的输出作为新模型的输入：

from tensorflow.keras.models import Model

# 截取输入到Flatten层的输出作为基模型，去掉原分类层
base_model = Model(
    inputs=pretrained_model.input,
    outputs=pretrained_model.get_layer('flatten').output
)

2.3 新增自定义层构建新模型 #

在基模型后添加你需要的自定义网络层，适配新的任务：

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout

new_model = Sequential()
# 添加预训练的特征提取基模型
new_model.add(base_model)
# 新增自定义层
new_model.add(Dense(128, activation='relu'))
new_model.add(Dropout(0.3))
# 替换为适配新任务的输出层，比如新任务是5分类就设置units=5
new_model.add(Dense(5, activation='softmax'))

# 编译新模型后即可训练
new_model.compile(
    loss='categorical_crossentropy',
    optimizer='adam',
    metrics=['acc']
)

# 若后续需要微调预训练层，可将base_model.trainable设为True，调小学习率再训练几轮即可

2.4 仅加载权重的操作步骤 #

如果你采用方案2仅保存了权重，需要先原封不动复现你原有的模型结构，再加载权重：

# 先复现原模型结构，和你之前写的代码完全一致
model_cnn = Sequential()
model_cnn.add(Embedding(vocab_size, embedding_size, input_length=max_len))
model_cnn.add(Conv1D(32, 3, padding='same', activation='relu'))
model_cnn.add(Conv1D(64, 3, padding='same', activation='relu'))
model_cnn.add(MaxPooling1D())
model_cnn.add(Flatten())
model_cnn.add(Dense(250, activation='relu'))
model_cnn.add(Dense(2, activation='softmax'))

# 加载预训练权重
model_cnn.load_weights('cnn_pretrained_weights.h5')

# 后续截取层、加新层的操作和上述步骤完全一致

内容的提问来源于stack exchange，提问作者Jugal Kishore Chanda