You need to enable JavaScript to run this app.
最新活动
大模型
产品
解决方案
定价
生态与合作
支持与服务
开发者
了解我们

如何在Keras中仅将模型与权重保存到单个文件?

精简Keras模型保存体积的实用方案

嘿,这个问题我之前帮不少开发者踩过坑——Keras默认的model.save()确实会把训练相关的所有细节都打包进去,几百MB的体积确实够头疼的。结合你提到的Keras FAQ里的存储内容,咱可以针对性地砍掉不需要的部分,下面给你几个实用的方案:

1. 拆分保存:只存架构+权重(最常用的精简方式)

如果不需要从断点恢复训练,完全可以把模型架构和权重分开保存,这样能砍掉训练配置和优化器状态这两块占空间的内容。

保存步骤:

# 保存模型架构为JSON文件(体积极小)
model_json = model.to_json()
with open("model_architecture.json", "w") as json_file:
    json_file.write(model_json)

# 单独保存权重文件(这是主要体积来源,但比完整模型小很多)
model.save_weights("model_weights.h5")

加载步骤:

from tensorflow.keras.models import model_from_json

# 加载架构
json_file = open('model_architecture.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
loaded_model = model_from_json(loaded_model_json)

# 加载权重
loaded_model.load_weights("model_weights.h5")

# 如果需要推理,记得编译(如果之前没编译过)
loaded_model.compile(optimizer='adam', loss='categorical_crossentropy')

这个方案的好处是架构文件几乎可以忽略不计,权重文件只保留模型参数,体积能砍到原来的一半甚至更少,而且管理起来也灵活。

2. 用save()时禁用优化器状态

如果你不想拆分文件,还想保留架构和训练配置,但不需要恢复训练,可以直接在保存时加上include_optimizer=False参数,这样就不会存储优化器的状态了:

# HDF5格式
model.save("my_slim_model.h5", include_optimizer=False)

# SavedModel格式(TensorFlow推荐格式)
model.save("my_slim_model", include_optimizer=False)

这个操作很简单,直接砍掉了优化器状态这块冗余内容,体积会明显缩小,而且加载模型的方式和原来完全一样,不需要额外步骤。

3. 模型量化:进一步压缩(适合部署场景)

如果还想把体积压得更小,比如要部署到移动端或边缘设备,可以试试模型量化——把权重从float32转换成float16甚至int8,体积直接砍半甚至减到原来的1/4,而且精度损失几乎可以忽略。

示例:转成TensorFlow Lite量化模型

import tensorflow as tf

# 转换为float16量化模型
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
tflite_model = converter.convert()

# 保存量化后的模型
with open("quantized_model.tflite", "wb") as f:
    f.write(tflite_model)

如果你的模型是用TensorFlow训练的,还可以用量化感知训练,在训练过程中就模拟量化效果,最终的量化模型精度会更高。

4. 导出推理专用模型(去掉训练专属层)

如果你的模型里有训练时才用到的层(比如DropoutBatchNormalization的训练模式),可以导出一个纯推理用的模型,不仅体积更小,推理速度也更快:

# 构建推理模型(自动切换到推理模式)
inference_model = tf.keras.models.Model(inputs=model.input, outputs=model.output)

# 保存推理模型
inference_model.save("inference_only_model.h5", include_optimizer=False)

这样保存的模型会自动把训练时的层切换到推理状态,去掉了训练相关的参数,体积进一步精简。


内容的提问来源于stack exchange,提问作者Adam Smith

火山引擎 最新活动