各位观众老爷，大家好！今天咱们不聊诗和远方，就聊聊怎么让你的浏览器变成一个“算命先生”（误），哦不，是机器学习专家！
今天的主题是 JS AI / ML with TensorFlow.js / ONNX Runtime Web：浏览器端机器学习部署与优化。 简单来说，就是如何在浏览器里玩转机器学习，而且要玩得溜，玩得快。
咱们先来个暖场小剧场：
- 你: 浏览器，你能不能给我预测一下明天的股票涨不涨？
- 浏览器: (吭哧吭哧算了一天) 我也不知道啊…我只是个浏览器啊！
- 你: (甩出一行 TensorFlow.js 代码) 现在呢？
- 浏览器: (两秒钟搞定) 大概率会涨！(信不信由你…)
是不是感觉很神奇？好，废话不多说，咱们正式开始！
第一幕：浏览器端的机器学习？凭什么？
可能有些小伙伴会觉得奇怪，机器学习不是应该在服务器上跑吗？浏览器这小身板，能行吗？答案是：能！而且好处多多！
- 隐私保护: 数据不出浏览器，用户隐私更有保障。毕竟，谁也不想自己的数据被别人扒光了。
- 降低服务器压力: 计算都在客户端完成，服务器压力大大减轻。想想双十一抢购，要是都靠服务器，那得瘫痪成啥样？
- 离线应用: 即使没有网络，也能进行一些简单的预测和分析。比如，离线版的“垃圾分类助手”，随时随地帮你分辨。
- 更快的响应速度: 省去了网络传输的延迟，响应速度更快。用户体验杠杠的！
第二幕：TensorFlow.js：老牌劲旅，简单易用
TensorFlow.js 顾名思义，就是 TensorFlow 在 JavaScript 领域的化身。它允许你在浏览器和 Node.js 环境中训练和部署机器学习模型。
2.1 Hello, TensorFlow.js!
咱们先来个最简单的例子，让浏览器认识一下 TensorFlow.js。
// 引入 TensorFlow.js 库
import * as tf from '@tensorflow/tfjs';
// 创建一个简单的模型：线性回归
const model = tf.sequential();
model.add(tf.layers.dense({units: 1, inputShape: [1]})); //units 输出维度，inputShape 输入维度
// 指定损失函数和优化器
model.compile({loss: 'meanSquaredError', optimizer: 'sgd'}); //sgd随机梯度下降
// 准备训练数据
const xs = tf.tensor2d([1, 2, 3, 4], [4, 1]); // 4个样本，每个样本1个特征
const ys = tf.tensor2d([2, 4, 6, 8], [4, 1]); // 对应的标签
// 训练模型
async function trainModel() {
await model.fit(xs, ys, {epochs: 100}); //epochs 训练轮数
console.log('模型训练完成！');
}
trainModel();
// 预测
async function predict(x) {
const tensorX = tf.tensor2d([x], [1, 1]);
const prediction = model.predict(tensorX);
prediction.print();
}
// 预测 x=5 的结果
predict(5);这段代码实现了一个简单的线性回归模型，用来预测 y = 2 * x。
代码解释:
- 引入库:
import * as tf from '@tensorflow/tfjs';引入 TensorFlow.js 库。 - 创建模型:
tf.sequential()创建一个序列模型，也就是一层一层堆叠起来的模型。tf.layers.dense()添加一个全连接层。units: 1表示输出维度为 1，inputShape: [1]表示输入维度为 1。 - 编译模型:
model.compile()配置模型的损失函数和优化器。loss: 'meanSquaredError'表示使用均方误差作为损失函数，optimizer: 'sgd'表示使用随机梯度下降作为优化器。 - 准备数据:
tf.tensor2d()创建二维张量。xs是输入数据，ys是标签数据。 - 训练模型:
model.fit()训练模型。epochs: 100表示训练 100 轮。 - 预测:
model.predict()使用训练好的模型进行预测。
2.2 加载预训练模型
自己训练模型当然很酷，但有时候我们更喜欢“拿来主义”，直接使用别人训练好的模型。TensorFlow.js 支持加载各种格式的预训练模型，例如 TensorFlow SavedModel、Keras 模型等。
// 加载预训练模型 (假设模型文件为 model.json 和 weights.bin)
async function loadModel() {
const model = await tf.loadLayersModel('model.json');
console.log('模型加载完成！');
return model;
}
// 使用模型进行预测
async function predict(model, inputData) {
const tensorInput = tf.tensor(inputData, [1, inputData.length]); // 假设输入数据是一维数组
const prediction = model.predict(tensorInput);
prediction.print();
}
// 主函数
async function main() {
const model = await loadModel();
const inputData = [0.1, 0.2, 0.3, 0.4]; // 示例输入数据
await predict(model, inputData);
}
main();代码解释:
- 加载模型:
tf.loadLayersModel()加载预训练模型。需要提供模型配置文件的路径 (model.json)。权重文件 (weights.bin) 通常会自动加载。 - 创建输入张量:
tf.tensor()将输入数据转换为张量。 - 预测:
model.predict()使用加载的模型进行预测。
2.3 迁移学习
迁移学习是指将一个模型在某个任务上学习到的知识迁移到另一个任务上。这可以大大减少训练时间和数据需求。
例如，我们可以使用预训练的 MobileNet 模型，将其最后一层替换为我们自己的分类器，然后用我们自己的数据进行微调。
// 加载预训练的 MobileNet 模型
async function loadMobileNet() {
const mobilenet = await tf.loadLayersModel(
'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json'); // MobileNet V1 模型地址
// 移除 MobileNet 的最后一层
const layer = mobilenet.getLayer('conv_preds');
const truncatedMobileNet = tf.model({inputs: mobilenet.inputs, outputs: layer.output});
truncatedMobileNet.trainable = false; // 冻结 MobileNet 的权重，防止微调时影响预训练的权重
return truncatedMobileNet;
}
// 创建自定义分类器
function createClassifier(numClasses) {
const model = tf.sequential();
model.add(tf.layers.flatten({inputShape: truncatedMobileNet.outputShape.slice(1)}));
model.add(tf.layers.dense({units: 128, activation: 'relu'}));
model.add(tf.layers.dense({units: numClasses, activation: 'softmax'}));
return model;
}
// 训练自定义分类器
async function trainClassifier(truncatedMobileNet, classifier, xs, ys) {
// 将 MobileNet 的输出作为分类器的输入
const embeddings = truncatedMobileNet.predict(xs);
classifier.compile({loss: 'categoricalCrossentropy', optimizer: 'adam', metrics: ['accuracy']});
await classifier.fit(embeddings, ys, {epochs: 10});
console.log('分类器训练完成！');
}
// 主函数
async function main() {
const truncatedMobileNet = await loadMobileNet();
const numClasses = 10; // 假设有 10 个类别
const classifier = createClassifier(numClasses);
// 准备训练数据 (xs 和 ys)
// xs: 输入图片张量 (例如，形状为 [numExamples, 224, 224, 3] 的张量)
// ys: 标签张量 (例如，形状为 [numExamples, numClasses] 的 one-hot 编码张量)
await trainClassifier(truncatedMobileNet, classifier, xs, ys);
// 使用训练好的模型进行预测
// ...
}
main();代码解释:
- 加载 MobileNet:
tf.loadLayersModel()加载预训练的 MobileNet 模型。 - 移除最后一层:
mobilenet.getLayer()获取 MobileNet 的最后一层，然后使用tf.model()创建一个新的模型，只包含 MobileNet 的中间层。 - 冻结权重:
truncatedMobileNet.trainable = false;冻结 MobileNet 的权重，防止在微调时影响预训练的权重。 - 创建分类器:
tf.sequential()创建一个新的序列模型，作为自定义分类器。 - 训练分类器:
truncatedMobileNet.predict()将输入图片通过 MobileNet 提取特征，然后将这些特征作为分类器的输入进行训练。
第三幕：ONNX Runtime Web：后起之秀，性能至上
ONNX (Open Neural Network Exchange) 是一种开放的模型表示格式，允许在不同的深度学习框架之间交换模型。ONNX Runtime 是一个高性能的 ONNX 模型推理引擎。ONNX Runtime Web 允许你在浏览器中运行 ONNX 模型，并且通常比 TensorFlow.js 具有更好的性能。
3.1 Hello, ONNX Runtime Web!
<!DOCTYPE html>
<html>
<head>
<title>ONNX Runtime Web Demo</title>
<script src="https://cdn.jsdelivr.net/npm/[email protected]/dist/ort.min.js"></script>
</head>
<body>
<script>
async function runONNX() {
// 加载 ONNX 模型 (假设模型文件为 model.onnx)
const session = await ort.InferenceSession.create('model.onnx');
// 准备输入数据
const inputTensor = new ort.Tensor('float32', new Float32Array([1.0, 2.0, 3.0]), [1, 3]); // 示例输入数据
// 运行推理
const feeds = {'input': inputTensor};
const results = await session.run(feeds);
// 获取输出
const outputTensor = results.output;
console.log('Output:', outputTensor.data);
}
runONNX();
</script>
</body>
</html>代码解释:
- 引入库:
<script src="https://cdn.jsdelivr.net/npm/[email protected]/dist/ort.min.js"></script>引入 ONNX Runtime Web 库。 - 加载模型:
ort.InferenceSession.create()加载 ONNX 模型。需要提供模型文件的路径 (model.onnx)。 - 创建输入张量:
new ort.Tensor()创建一个 ONNX 张量。需要指定数据类型 (float32)、数据 (new Float32Array([1.0, 2.0, 3.0])) 和形状 ([1, 3])。 - 运行推理:
session.run()运行推理。需要提供一个包含输入张量的对象 (feeds)。 - 获取输出:
results.output获取输出张量。
3.2 ONNX 模型转换
想要在 ONNX Runtime Web 中运行模型，首先需要将模型转换为 ONNX 格式。可以使用各种工具进行转换，例如 TensorFlow、PyTorch 等。
-
TensorFlow to ONNX: 可以使用
tf2onnx工具将 TensorFlow 模型转换为 ONNX 格式。python -m tf2onnx.convert --saved-model saved_model_dir --output model.onnx --opset 13 -
PyTorch to ONNX: 可以使用
torch.onnx.export()函数将 PyTorch 模型转换为 ONNX 格式。import torch # 创建一个简单的 PyTorch 模型 class MyModel(torch.nn.Module): def __init__(self): super(MyModel, self).__init__() self.linear = torch.nn.Linear(3, 2) def forward(self, x): return self.linear(x) model = MyModel() # 创建一个示例输入 dummy_input = torch.randn(1, 3) # 导出模型到 ONNX 格式 torch.onnx.export(model, dummy_input, 'model.onnx', verbose=True, input_names=['input'], output_names=['output'])
第四幕：性能优化：让你的浏览器飞起来
即使使用了 ONNX Runtime Web，也并不意味着你的模型就能跑得飞快。还需要进行一些性能优化，才能让你的浏览器真正“飞起来”。
4.1 模型量化
模型量化是指将模型的权重从浮点数转换为整数。这可以大大减少模型的大小，并提高推理速度。
- TensorFlow.js: 可以使用 TensorFlow Lite 的量化工具将 TensorFlow.js 模型量化。
- ONNX Runtime Web: ONNX Runtime 支持各种量化方案，例如动态量化、静态量化等。
4.2 模型剪枝
模型剪枝是指移除模型中不重要的连接或神经元。这可以减少模型的复杂度，并提高推理速度。
- TensorFlow.js: 可以使用 TensorFlow Model Optimization Toolkit 进行模型剪枝。
- ONNX Runtime Web: ONNX Runtime 支持使用 SparseTensor API 进行模型剪枝。
4.3 WebAssembly (WASM)
WebAssembly 是一种新的二进制指令格式，可以在现代浏览器中以接近原生速度运行。使用 WASM 可以大大提高 JavaScript 代码的性能。
- TensorFlow.js: TensorFlow.js 支持使用 WASM 后端，可以提高模型推理速度。
- ONNX Runtime Web: ONNX Runtime Web 默认使用 WASM 后端。
4.4 WebGL
WebGL 是一种 JavaScript API，用于在浏览器中渲染 2D 和 3D 图形。使用 WebGL 可以将一些计算任务卸载到 GPU 上，从而提高性能。
- TensorFlow.js: TensorFlow.js 支持使用 WebGL 后端，可以加速张量运算。
- ONNX Runtime Web: ONNX Runtime Web 也支持使用 WebGL 后端。
4.5 其他优化技巧
- 避免不必要的内存分配: 频繁的内存分配会导致性能下降。尽量重用张量，避免创建过多的临时变量。
- 使用异步操作: 避免阻塞主线程。使用
async/await进行异步操作，防止 UI 卡顿。 - 优化数据预处理: 数据预处理是机器学习流程中重要的一环。优化数据预处理流程可以提高整体性能。
第五幕：总结与展望
今天我们一起探索了如何在浏览器端部署和优化机器学习模型。从 TensorFlow.js 到 ONNX Runtime Web，从模型量化到 WebAssembly，我们学习了各种技术和技巧。
| 技术 | 优点 | 缺点 | 适用场景 |
|---|---|---|---|
| TensorFlow.js | 简单易用，生态完善，支持模型训练，适合快速原型开发 | 性能相对较弱，尤其是在移动设备上 | 简单的模型，对性能要求不高的场景，需要在线训练的场景 |
| ONNX Runtime Web | 性能优异，尤其是在大型模型上，支持多种量化方案 | 学习曲线较陡峭，生态不如 TensorFlow.js 完善，不支持模型训练 | 对性能要求较高的场景，需要运行大型预训练模型的场景 |
| 模型量化 | 减少模型大小，提高推理速度 | 可能会降低模型精度 | 所有场景，尤其是在资源受限的设备上 |
| 模型剪枝 | 减少模型复杂度，提高推理速度 | 可能会降低模型精度 | 所有场景，尤其是在模型过于庞大的情况下 |
| WebAssembly (WASM) | 提供接近原生速度的性能 | 调试困难 | 所有场景，尤其是在需要高性能的情况下 |
| WebGL | 将计算任务卸载到 GPU 上，提高性能 | 依赖 GPU 性能，兼容性问题 | 计算密集型任务，例如图像处理、视频分析等 |
未来，浏览器端的机器学习将会越来越普及。随着 WebAssembly 和 WebGPU 等技术的不断发展，浏览器端的计算能力将会越来越强大。我们可以期待更多的 AI 应用在浏览器中涌现，为用户带来更智能、更便捷的体验。
好了，今天的讲座就到这里。希望大家有所收获！如果有什么问题，欢迎在评论区留言。
咱们下次再见！(挥手)