TensorFlow 模型优化工具包

TensorFlow 模型优化工具包是一套工具，无论初学者还是高级用户都可以使用它来优化机器学习模型，以便部署和执行。

支持的技术包括量化和稀疏权重剪枝。此外，还为 Keras 提供了专门的 API。

有关该项目及其各个工具的概述、优化收益以及我们的路线图，请参阅 tensorflow.org/model_optimization。 该网站还提供了各种教程和 API 文档。

该工具包提供稳定的 Python API。

安装

有关安装说明，请参阅 tensorflow.org/model_optimization/guide/install。

贡献指南

如果您希望为 TensorFlow 模型优化项目做出贡献，请务必阅读 贡献指南。本项目遵守 TensorFlow 的 行为准则。 参与时，您应遵守该准则。

我们使用 GitHub 问题 来跟踪请求和错误。

维护者

子包	维护者
tfmot.clustering	Arm ML 工具链团队
tfmot.quantization	TensorFlow 模型优化团队
tfmot.sparsity	TensorFlow 模型优化团队

社区

作为 TensorFlow 的一部分，我们致力于营造开放且友好的环境。

• TensorFlow 博客：及时了解 TensorFlow 团队的最新内容以及社区精选文章。

TensorFlow Model Optimization Toolkit 快速上手指南

TensorFlow Model Optimization Toolkit 是一套用于优化机器学习模型以便部署和执行的工具套件，支持量化（Quantization）和剪枝（Pruning）等技术，并提供专为 Keras 设计的 API。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：3.7 - 3.10
• 前置依赖：
  • TensorFlow 2.x (推荐最新稳定版)
  • NumPy
• 网络建议：国内开发者建议使用清华或阿里镜像源加速安装。

安装步骤

方法一：使用 pip 安装（推荐）

使用国内镜像源加速安装：

pip install tensorflow-model-optimization -i https://pypi.tuna.tsinghua.edu.cn/simple

或使用官方源：

pip install tensorflow-model-optimization

方法二：从源码安装（可选）

如需使用最新开发版功能：

git clone https://github.com/tensorflow/model-optimization.git
cd model-optimization
pip install .

基本使用

以下示例展示如何使用该工具包对 Keras 模型进行后训练量化（Post-training Quantization），这是最基础且常用的优化场景。

1. 导入依赖并加载模型

import tensorflow as tf
import tensorflow_model_optimization as tfmot

# 加载预训练的 Keras 模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')

2. 定义量化配置并应用

# 创建默认的量化工具配置
quantize_config = tfmot.quantization.keras.QuantizeConfig()

# 将量化包装器应用到模型
quantized_model = tfmot.quantization.keras.quantize_model(model)

3. 编译与训练（微调）

量化后的模型需要重新编译。通常需要进行少量的微调以恢复精度。

quantized_model.compile(
    optimizer='adam',
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

# 执行微调训练 (示例数据需替换为实际数据集)
# quantized_model.fit(x_train, y_train, epochs=5)

4. 导出模型

训练完成后，移除量化包装器并保存最终的 TFLite 或 SavedModel 格式。

# 移除量化包装器，获取最终模型
final_model = tfmot.quantization.keras.strip_quantize(quantized_model)

# 转换为 TFLite 格式以获得最佳移动端性能
converter = tf.lite.TFLiteConverter.from_keras_model(final_model)
tflite_model = converter.convert()

with open('model_quantized.tflite', 'wb') as f:
    f.write(tflite_model)

提示：除了量化，该工具包还支持 tfmot.sparsity 进行模型剪枝，以及 tfmot.clustering 进行权重聚类，具体用法可参考官方文档。