JS `Edge AI` (`TensorFlow.js` / `ONNX Runtime Web`) `Model Quantization` 与 `Pruning`

嘿，大家好！我是今天的主讲人，很高兴能和大家一起聊聊如何在浏览器里“榨干”AI模型的最后一点性能！今天要讲的是JS Edge AI (也就是 TensorFlow.js / ONNX Runtime Web) 中的模型量化和剪枝，让咱们的AI模型在前端也能跑得飞起。

开场白：为啥要在浏览器里搞事情？

先来说说为啥我们要费劲巴拉地在浏览器里搞AI。原因很简单：隐私！数据不出门，安全又放心。想象一下，用户上传一张照片，你想识别里面的物体，如果把照片传到服务器，再返回结果，速度慢不说，用户隐私也暴露了。但在浏览器里直接跑，速度快，隐私有保障，简直完美！

但是！问题来了，在浏览器里跑AI模型，资源有限，性能受限。大型模型跑起来慢不说，还耗电，简直是移动设备的噩梦。这时候，模型量化和剪枝就派上用场了，它们就像是给模型做了个“瘦身”，让它跑得更快，更省资源。

第一部分：模型量化 (Quantization) — 压缩模型的数值精度

模型量化，顾名思义，就是把模型里的数值精度降低。通常，深度学习模型使用32位浮点数 (float32) 来表示权重和激活值。量化的目的就是把这些32位的“胖子”变成8位的“瘦子”，甚至更小的“小矮人”。

1.1 啥是浮点数？为啥要量化？

咱们先简单回顾一下浮点数。浮点数是一种表示实数的方法，它用一部分位数表示整数部分，另一部分位数表示小数部分。32位浮点数能表示很大的数值范围，精度也很高。

但是！精度越高，占用的存储空间就越大，计算速度也越慢。想象一下，你要计算1.23456789 9.87654321，和计算1.2 9.9，哪个更快？当然是后者！

所以，量化的本质就是牺牲一定的精度，来换取更小的模型体积和更快的计算速度。

1.2 量化的种类

量化方法有很多种，但常见的有以下几种：

• 动态量化 (Dynamic Quantization): 在运行时确定量化参数，比如最大值和最小值。
• 静态量化 (Static Quantization): 在训练后，使用校准数据集确定量化参数。
• 训练时量化 (Quantization Aware Training, QAT): 在训练过程中模拟量化，让模型适应量化带来的影响。

1.3 TensorFlow.js 中的量化

TensorFlow.js 提供了模型量化的工具，但目前支持度还不够完善。不过，我们可以使用 TensorFlow Python 来量化模型，然后转换成 TensorFlow.js 模型。

1.3.1 TensorFlow Python 量化示例

import tensorflow as tf

# 加载预训练模型
model = tf.keras.models.load_model('my_model.h5')

# 量化配置
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

# (可选) 如果使用静态量化，需要提供校准数据集
# def representative_dataset():
#   for _ in range(100):
#     data = np.random.rand(1, 224, 224, 3).astype(np.float32)
#     yield [data]
# converter.representative_dataset = representative_dataset
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# converter.inference_input_type = tf.uint8
# converter.inference_output_type = tf.uint8

# 转换成 TFLite 模型
tflite_model = converter.convert()

# 保存 TFLite 模型
with open('my_model.tflite', 'wb') as f:
  f.write(tflite_model)

这段代码做了什么？

1. 加载模型: tf.keras.models.load_model('my_model.h5') 加载你的 Keras 模型。
2. 创建 Converter: tf.lite.TFLiteConverter.from_keras_model(model) 创建一个转换器，用于将 Keras 模型转换为 TFLite 模型。
3. 设置优化: converter.optimizations = [tf.lite.Optimize.DEFAULT] 开启默认的优化，其中就包括量化。
4. (可选) 设置校准数据集: 如果使用静态量化，你需要提供一个校准数据集，让转换器确定量化参数。
5. 转换模型: tflite_model = converter.convert() 将模型转换为 TFLite 模型。
6. 保存模型: with open('my_model.tflite', 'wb') as f: f.write(tflite_model) 将 TFLite 模型保存到文件。

1.3.2 在 TensorFlow.js 中加载 TFLite 模型

import * as tf from '@tensorflow/tfjs';
import * as tflite from '@tensorflow/tfjs-tflite';

async function loadModel() {
  const model = await tflite.loadTFLiteModel('my_model.tflite');
  return model;
}

async function predict(model, inputTensor) {
  const outputTensor = model.predict(inputTensor);
  return outputTensor;
}

// 示例用法
async function runInference() {
  const model = await loadModel();
  // 创建一个输入张量
  const inputTensor = tf.tensor([ /* 你的输入数据 */ ]);
  const outputTensor = await predict(model, inputTensor);
  // 处理输出张量
  console.log(outputTensor.dataSync());
}

runInference();

这段代码展示了如何在 TensorFlow.js 中加载和使用 TFLite 模型。

1.4 ONNX Runtime Web 中的量化

ONNX Runtime Web 也支持量化模型，但需要先将模型转换为 ONNX 格式，然后再进行量化。

1.4.1 ONNX 量化示例 (使用 ONNX 提供的工具)

import onnx
from onnxruntime.quantization import quantize_dynamic, quantize_static, CalibrationDataReader, QuantType

# 加载 ONNX 模型
model_path = "my_model.onnx"
quantized_model_path = "my_model_quantized.onnx"

# 动态量化
quantize_dynamic(model_path, quantized_model_path, weight_type=QuantType.QUInt8)

# 静态量化 (需要校准数据集)
# class MyDataReader(CalibrationDataReader):
#   def __init__(self, calibration_data_path):
#     self.calibration_data_path = calibration_data_path
#     self.enum_data = None

#   def get_next(self):
#     if self.enum_data is None:
#       self.enum_data = iter(...) # 加载你的校准数据集
#     try:
#       return next(self.enum_data)
#     except StopIteration:
#       return None

# data_reader = MyDataReader(...)
# quantize_static(model_path, quantized_model_path, data_reader, quant_format=QuantFormat.QDQ, weight_type=QuantType.QUInt8)

这段代码展示了如何使用 ONNX 提供的工具进行模型量化。动态量化不需要校准数据集，静态量化则需要。

1.4.2 在 ONNX Runtime Web 中加载 ONNX 模型

import { InferenceSession } from 'onnxruntime-web';

async function loadModel() {
  const session = await InferenceSession.create('my_model_quantized.onnx');
  return session;
}

async function predict(session, inputTensor) {
  const feeds = { 'input': inputTensor }; // 'input' 是输入张量的名称
  const results = await session.run(feeds);
  return results;
}

// 示例用法
async function runInference() {
  const session = await loadModel();
  // 创建一个输入张量
  const inputTensor = new onnx.Tensor(new Float32Array([ /* 你的输入数据 */ ]), 'float32', [ /* 输入张量的维度 */ ]);
  const results = await predict(session, inputTensor);
  // 处理输出张量
  console.log(results);
}

runInference();

这段代码展示了如何在 ONNX Runtime Web 中加载和使用量化后的 ONNX 模型。

1.5 量化的注意事项

• 精度损失: 量化会降低模型的精度，需要在精度和性能之间找到平衡。
• 校准数据集: 静态量化需要校准数据集，数据集的质量会影响量化效果。
• 硬件支持: 不同的硬件对量化的支持程度不同，需要根据实际情况选择合适的量化方法。

第二部分：模型剪枝 (Pruning) — 移除模型中不重要的连接

模型剪枝，简单来说，就是把模型里不重要的连接 (权重) 砍掉，就像给一棵树修剪枝叶一样。这样可以减少模型的参数数量，降低计算复杂度，提高推理速度。

2.1 啥是不重要的连接？

模型中的连接，也就是权重，有些连接对模型的预测结果影响很大，有些连接则几乎没有影响。我们可以把那些影响小的连接认为是“不重要的连接”。

2.2 剪枝的种类

常见的剪枝方法有以下几种：

• 非结构化剪枝 (Unstructured Pruning): 随机地移除模型中的权重。
• 结构化剪枝 (Structured Pruning): 移除整个神经元或卷积核。

2.3 TensorFlow.js 中的剪枝

TensorFlow.js 本身没有提供剪枝的工具，但我们可以使用 TensorFlow Python 来剪枝模型，然后转换成 TensorFlow.js 模型。

2.3.1 TensorFlow Python 剪枝示例

import tensorflow as tf
import tensorflow_model_optimization as tfmot

# 加载预训练模型
model = tf.keras.models.load_model('my_model.h5')

# 定义剪枝参数
pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
          initial_sparsity=0.50,
          final_sparsity=0.90,
          begin_step=0,
          end_step=1000
      )
}

# 应用剪枝
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, **pruning_params)

# 编译模型
model_for_pruning.compile(optimizer='adam',
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

# 训练模型
logdir = tempfile.mkdtemp()
callbacks = [
  tfmot.sparsity.keras.UpdatePruning(),
  tfmot.sparsity.keras.PruningSummaries(log_dir=logdir, profile_model=False)
]

model_for_pruning.fit(x_train, y_train,
                  epochs=10,
                  callbacks=callbacks)

# 去除剪枝层
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

# 保存模型
model_for_export.save('my_model_pruned.h5')

这段代码做了什么？

1. 加载模型: tf.keras.models.load_model('my_model.h5') 加载你的 Keras 模型。
2. 定义剪枝参数: pruning_params 定义了剪枝的参数，包括初始稀疏度、最终稀疏度、开始步数和结束步数。
3. 应用剪枝: tfmot.sparsity.keras.prune_low_magnitude(model, **pruning_params) 将剪枝应用到模型上。
4. 编译模型: model_for_pruning.compile(...) 编译模型。
5. 训练模型: model_for_pruning.fit(...) 训练模型，注意要使用 tfmot.sparsity.keras.UpdatePruning() 回调函数来更新剪枝。
6. 去除剪枝层: tfmot.sparsity.keras.strip_pruning(model_for_pruning) 去除剪枝层，得到最终的剪枝模型。
7. 保存模型: model_for_export.save('my_model_pruned.h5') 保存剪枝后的模型。

2.4 ONNX Runtime Web 中的剪枝

ONNX Runtime Web 同样没有直接提供剪枝的工具，需要先在其他框架中剪枝，然后转换为 ONNX 格式。

2.5 剪枝的注意事项

• 精度损失: 剪枝会降低模型的精度，需要在精度和性能之间找到平衡。
• 剪枝率: 剪枝率越高，模型体积越小，但精度损失也越大。
• 迭代剪枝: 可以采用迭代剪枝的方法，逐步提高剪枝率，并在每次剪枝后进行微调，以减少精度损失。

第三部分：量化 + 剪枝 = 性能起飞

通常情况下，我们会将量化和剪枝结合起来使用，以达到最佳的性能优化效果。先剪枝，去除不重要的连接，再量化，降低数值精度，双管齐下，让模型在浏览器里也能跑得飞起。

3.1 示例流程

1. 训练一个高精度模型。
2. 对模型进行剪枝。
3. 对剪枝后的模型进行量化。
4. 将量化后的模型转换为 TensorFlow.js 或 ONNX Runtime Web 支持的格式。
5. 在浏览器中加载和使用模型。

第四部分：实战案例 (伪代码)

假设我们有一个图像分类模型，需要在浏览器里实现实时分类。

// 1. 加载模型 (量化 + 剪枝)
const model = await tflite.loadTFLiteModel('my_model_quantized_pruned.tflite');

// 2. 获取摄像头图像
const video = document.getElementById('webcam');
const canvas = document.getElementById('canvas');
const ctx = canvas.getContext('2d');

async function startWebcam() {
  const stream = await navigator.mediaDevices.getUserMedia({ video: true });
  video.srcObject = stream;
  await video.play();
  canvas.width = video.videoWidth;
  canvas.height = video.videoHeight;
}

// 3. 图像预处理
function preprocessImage(image) {
  ctx.drawImage(image, 0, 0, canvas.width, canvas.height);
  const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
  const tensor = tf.browser.fromPixels(imageData)
    .resizeNearestNeighbor([224, 224])
    .toFloat()
    .div(tf.scalar(255.0))
    .expandDims();
  return tensor;
}

// 4. 模型推理
async function predict() {
  const tensor = preprocessImage(video);
  const predictions = await model.predict(tensor).data();
  // 5. 处理结果 (例如显示分类结果)
  console.log(predictions);
  tensor.dispose(); // 释放内存
  requestAnimationFrame(predict); // 循环推理
}

// 启动摄像头和推理
await startWebcam();
predict();

这个伪代码展示了一个简单的图像分类流程，包括加载量化和剪枝后的模型、获取摄像头图像、图像预处理、模型推理和结果处理。

第五部分：总结与展望

今天我们聊了如何在 JS Edge AI 中使用模型量化和剪枝来优化模型性能。量化和剪枝是两种常用的模型压缩技术，它们可以显著减小模型体积，提高推理速度，让我们的AI模型在浏览器里也能跑得飞起。

技术	优点	缺点	适用场景
量化	减小模型体积，提高推理速度，降低功耗	可能降低模型精度，需要校准数据集 (静态量化)	资源受限的设备 (如移动设备)，对精度要求不高的场景
剪枝	减小模型体积，提高推理速度，降低计算复杂度	可能降低模型精度，需要重新训练或微调模型	大型模型，计算资源有限的场景
量化+剪枝	结合了量化和剪枝的优点，可以达到最佳的性能优化效果	需要权衡精度和性能，需要更多的调参和实验	对性能要求极高，同时对精度有一定要求的场景

未来，随着硬件的不断发展和算法的不断优化，JS Edge AI 的应用场景将会越来越广泛。我们可以期待更多的AI应用在浏览器里落地生根，为用户带来更好的体验。

最后，感谢大家的聆听！ 希望今天的分享对大家有所帮助。如果有什么问题，欢迎随时提问！