Featured image of post 使用 TensorFlow.js 在浏览器中进行机器学习Featured image of post 使用 TensorFlow.js 在浏览器中进行机器学习

使用 TensorFlow.js 在浏览器中进行机器学习

使用 TensorFlow.js 在浏览器中进行机器学习。涵盖模型训练、迁移学习、WebGL/WebGPU 后端、姿势检测和隐私保护 ML。

浏览器中的机器学习消除了服务器成本,保护用户隐私,并支持离线智能应用程序。 TensorFlow.js 通过由 WebGL 和 WebGPU 后端提供支持的 GPU 加速推理和训练,为 JavaScript 开发人员带来了 ML。本文涵盖加载预训练模型、迁移学习、实时姿态检测和生产部署注意事项。

为什么要在浏览器中进行机器学习?

在客户端运行 ML 模型具有四个关键优势:零服务器成本(推理在用户设备上运行)、完全隐私(数据永远不会离开机器)、离线功能(模型加载后无需网络)和低延迟(无需往返预测)。权衡包括有限的计算能力、内存限制、移动设备的电池消耗以及大模型下载大小(5-20​​0 MB)。


TensorFlow.js 概述

TensorFlow.js 提供了三个 API。 核心 API 提供低级张量运算。 Layers API 支持类似于 Keras 的高级模型构建。 转换器 将 Python 训练的 Keras 或 TF 模型转换为 TensorFlow.js 格式:

tensorflowjs_converter --input_format=keras \
  path/to/model.h5 path/to/tfjs-model/

GPU 加速可通过 WebGL(成熟、广泛支持、仅限 float32)和 WebGPU(更好的性能、float16 支持、Chrome 113+)实现。 WASM 后端使用 XNNPACK 优化作为基于 CPU 的后备。


模型加载和推理

加载预训练模型并运行推理仅需要几行代码:

const model = await tf.loadGraphModel(
  'https://tfhub.dev/google/tfjs-model/mobilenet_v2/1/default/1/model.json'
);
const img = document.getElementById('image');
const tensor = tf.browser.fromPixels(img).expandDims(0);
const predictions = await model.predict(tensor);
const topClass = predictions.argMax(1).dataSync()[0];

内存管理至关重要。使用 tf.tidy() 自动处理中间张量:

const result = tf.tidy(() => {
  const processed = imgTensor.div(255).expandDims(0);
  return model.predict(processed);
});

迁移学习和自定义模型

迁移学习可以使用小数据集训练自定义分类器。该过程加载预训练的特征提取器(如 MobileNet),删除其顶部分类层,添加新的可训练层,并冻结基础层:

const baseModel = await tf.loadGraphModel(MOBILENET_URL);
const classifier = tf.sequential();
classifier.add(tf.layers.dense({ units: 128, activation: 'relu', inputShape: [1024] }));
classifier.add(tf.layers.dense({ units: NUM_CLASSES, activation: 'softmax' }));

// Freeze base model and train only the classifier
baseModel.trainable = false;
const combined = tf.sequential();
combined.add(baseModel);
combined.add(classifier);
combined.compile({ optimizer: 'adam', loss: 'categoricalCrossentropy' });
await combined.fit(trainingData, trainingLabels, { epochs: 10 });

此方法非常适合每个类只有 100 个样本的自定义图像分类任务。


实时姿势检测

TensorFlow.js 包含一个带有预构建检测器的模型动物园,用于姿势估计、对象检测和人脸地标:

const detector = await poseDetection.createDetector(
  poseDetection.SupportedModels.MoveNet
);
const video = document.getElementById('webcam');
const poses = await detector.estimatePoses(video);

// Draw keypoints on canvas overlay
poses.forEach(pose => {
  pose.keypoints.forEach(kp => {
    ctx.beginPath();
    ctx.arc(kp.x, kp.y, 5, 0, 2 * Math.PI);
    ctx.fillStyle = 'red';
    ctx.fill();
  });
});

性能优化技术包括跳帧、将输入调整为更小的尺寸以及选择更轻的模型变体。


后端性能比较

后端需要 GPU精密成熟度移动网络速度
网页GL是的浮动32成熟〜15ms
网络GPU是的浮动16/32实验〜8毫秒
WASM没有浮动32稳定〜40ms
中央处理器没有浮动32稳定〜120ms

对于常见模型,WebGPU 比 WebGL 提供大约 2 倍的加速,但浏览器支持目前仅限于基于 Chromium 的浏览器。


生产注意事项

模型可以是 5-200 MB。应用量化和修剪来减小尺寸。通过 tf.io.IndexedDB 在 IndexedDB 中缓存模型以避免重新下载。使用 tf.io 进度回调显示加载进度。在移动设备上,警告用户电池影响。认真使用 tf.tidy() 来防止内存泄漏。当硬件加速不可用时,始终提供优雅的回退。


隐私保护机器学习

客户端推理可确保敏感数据永远不会离开用户的设备。这对于医疗诊断支持、个人财务分析以及包含机密内容的文档分类尤其有价值。联邦学习和差分隐私概念可以应用于基于浏览器的训练场景。

TensorFlow.js 使 Web 开发人员可以使用 ML。从预先训练的模型和迁移学习开始,然后再尝试从头开始训练。为您的特定模型分析 WebGL 与 WebGPU 后端,并探索用于姿势检测、面部标志、手部跟踪和文本毒性分类的 tfjs-models 集合。