MobileNetV2基础模型剪枝后模型体积未变问题咨询

这种剪枝思路是可行的，但你没观察到模型体积缩小，核心原因是缺少了剪枝训练（权重稀疏化）和导出时的压缩优化两个关键环节，同时剪枝的基础配置也不够完整。

问题根源分析 #

1. 仅调用prune_low_magnitude只是给base_model套了剪枝包装器，此时模型权重还没有被稀疏化（大量权重变为0），自然不会有体积变化。
2. 即便完成剪枝训练，直接保存的Keras模型会保留所有权重（包括0值），需要移除剪枝包装器并配合压缩工具，才能体现出体积优势。

正确的剪枝流程 #

1. 配置剪枝参数并包装基础模型 #

先定义剪枝规则（比如目标稀疏度、剪枝起始步数），再包装base_model：

import tensorflow_model_optimization as tfmot

# 配置剪枝参数：训练1000步后，将90%的权重置为0
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.ConstantSparsity(
        target_sparsity=0.9,
        begin_step=1000,
        frequency=100
    ),
    'block_size': (1, 1),
    'block_pooling_type': 'AVG'
}

# 对base_model应用剪枝包装
pruned_base_model = tfmot.sparsity.keras.prune_low_magnitude(base_model, **pruning_params)

2. 组装模型并执行剪枝训练 #

组装完整模型后，必须通过训练让剪枝算法生效，把大量权重稀疏化为0：

# 组装带分类头的完整模型
model = tf.keras.models.Sequential([
    pruned_base_model,
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(
        units=10,
        activation='softmax',
        kernel_regularizer=tf.keras.regularizers.l2(0.01)
    )
])

# 编译模型，加入剪枝专用回调
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 剪枝回调：训练时更新剪枝状态，记录日志
callbacks = [
    tfmot.sparsity.keras.UpdatePruningStep(),
    tfmot.sparsity.keras.PruningSummaries(log_dir='./prune_logs')
]

# 用你的训练数据执行训练
model.fit(
    train_data,
    epochs=20,
    validation_data=val_data,
    callbacks=callbacks
)

3. 移除剪枝包装器并导出优化模型 #

训练完成后，移除剪枝相关的辅助组件，再通过压缩工具实现体积缩小：

# 移除剪枝包装器，得到仅含稀疏权重的模型
stripped_model = tfmot.sparsity.keras.strip_pruning(model)

# 保存稀疏模型
tf.keras.models.save_model(stripped_model, './pruned_mobilenetv2.h5')

# 可选：转换为TensorFlow Lite模型，进一步压缩体积
converter = tf.lite.TFLiteConverter.from_keras_model(stripped_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_pruned_model = converter.convert()

with open('./pruned_mobilenetv2.tflite', 'wb') as f:
    f.write(tflite_pruned_model)

额外注意事项 #

• 你的代码里仅剪枝base_model、保留分类头的逻辑是对的，分类头不会被剪枝。
• MobileNetV2中的BatchNormalization等层不可剪枝，工具会自动跳过，属于正常情况。
• 未压缩的稀疏模型体积和原模型差别不大，因为0值权重仍会被存储；只有通过TFLite优化或gzip压缩，才能体现出明显的体积缩减（通常能压缩到原模型的30%-50%）。

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