TensorFlow.js 前端机器学习：从模型加载到推理的完整指南

TensorFlow.js 前端机器学习：从模型加载到推理的完整指南

"在浏览器里跑 AI 模型"听起来像黑科技，但 TensorFlow.js 已经让它变成现实。

更重要的是，这不只是为了"酷"。前端 AI 带来了真实的商业价值：

• 隐私：用户数据不用发到服务器
• 速度：推理延迟低（毫秒级）
• 产品体验：即时反馈（图片分类、手势识别）

本文讲清楚怎样实战用 TensorFlow.js，避免"装库后不知道怎么用"。

1. TensorFlow.js 能做什么？

常见应用：

关键是，这些模型都是预训练的，你不需要自己训练，直接加载用。

2. 快速开始：图像分类

最常见的例子，用 MobileNet 分类图片：

<input type="file" id="imageFile" />
<img id="preview" style="max-width: 300px;" />
<p id="result"></p>

import * as tf from '@tensorflow/tfjs'
import * as mobilenet from '@tensorflow-models/mobilenet'

let model

async function loadModel() {
  model = await mobilenet.load()
  console.log('Model loaded')
}

document.getElementById('imageFile').addEventListener('change', async (e) => {
  const file = e.target.files[0]
  const img = document.getElementById('preview')
  img.src = URL.createObjectURL(file)
  
  img.onload = async () => {
    // 推理
    const predictions = await model.classify(img)
    
    const result = predictions
      .map(p => `${p.className}: ${(p.probability * 100).toFixed(2)}%`)
      .join('<br>')
    
    document.getElementById('result').innerHTML = result
  }
})

loadModel()

就这几行代码，你就有了一个"图片识别器"。

3. 模型格式与加载策略

TensorFlow.js 支持多种模型格式：

TFJS 格式（推荐）

const model = await tf.loadLayersModel(
  'file://.../model.json' // 或 HTTP URL
)

SavedModel（TF2 模型）

const model = await tf.loadGraphModel(
  'file://.../model.json'
)

ONNX 模型

const ort = require('onnxruntime-web')
const session = await ort.InferenceSession.create('model.onnx')

加载策略：

• 小模型（<50MB）：直接在 HTML 中 loadLayersModel
• 大模型：用 IndexedDB 本地缓存

const model = await tf.loadLayersModel(
  tf.io.indexedDB('my-model')
)

4. 推理优化：让模型跑得更快

问题：模型推理可能很慢（尤其是大模型）。

解决方案：

方案 A：使用更轻的模型

MobileNet is already optimized for mobile. For heavier tasks:

• PoseNet → BlazePose（更轻）
• ResNet → MobileNetV3（更轻）

方案 B：减小输入分辨率

// 不要用 2048x2048 的图片做推理
const resized = tf.image.resizeBilinear(img, [224, 224])
const predictions = await model.classify(resized)

方案 C：用 GPU（WebGL）

// 自动启用 GPU 加速（如果可用）
await tf.setBackend('webgl')
const predictions = await model.classify(img)

TensorFlow.js 支持多个后端：webgl、wasm、cpu。GPU 通常快 10x。

5. 实战：目标检测（COCO-SSD）

import * as cocoSsd from '@tensorflow-models/coco-ssd'

async function detectObjects(imageElement) {
  const model = await cocoSsd.load()
  
  const predictions = await model.estimateObjects(imageElement)
  
  // predictions = [
  //   { class: 'person', score: 0.96, bbox: [x, y, w, h] },
  //   { class: 'dog', score: 0.92, bbox: [...] }
  // ]
  
  const canvas = document.createElement('canvas')
  canvas.width = imageElement.width
  canvas.height = imageElement.height
  
  const ctx = canvas.getContext('2d')
  ctx.drawImage(imageElement, 0, 0)
  
  // 画检测框
  predictions.forEach(pred => {
    const [x, y, w, h] = pred.bbox
    ctx.strokeStyle = 'red'
    ctx.strokeRect(x, y, w, h)
    ctx.fillText(`${pred.class} ${pred.score.toFixed(2)}`, x, y - 5)
  })
  
  document.body.appendChild(canvas)
}

6. 内存管理与内存泄漏

TensorFlow.js 在浏览器里消耗内存很快。一个常见的坑是忘记清理张量：

// ❌ 内存泄漏
async function processImages(imageArray) {
  for (const img of imageArray) {
    const tensor = tf.browser.fromPixels(img)
    const result = model.predict(tensor) // tensor 没有被释放！
  }
}

// ✅ 正确做法
async function processImages(imageArray) {
  for (const img of imageArray) {
    const result = tf.tidy(() => {
      const tensor = tf.browser.fromPixels(img)
      return model.predict(tensor) // 自动清理中间张量
    })
  }
}

tf.tidy() 会自动清理其内部创建的所有中间张量，只保留最后的 result。

7. 隐私与数据安全

AI 在浏览器运行的最大优势是隐私：

// 用户的图片从不上传到服务器
const imageData = canvas.toDataURL() // 只保留在本地
const predictions = await model.classify(imageData)

// 只有分类结果可能被发送（不是原始图像）
analytics.track('image_classified', {
  classResult: predictions[0].className
})

这对文医疗、个人数据敏感的应用特别重要。

8. 常见陷阱与解决方案

陷阱 1：模型太大，加载很慢

解决：

• 缓存到 IndexedDB
• 在后台 Worker 加载模型
• 显示加载进度条

陷阱 2：推理时主线程卡顿

// ❌ 主线程卡
const predictions = await model.predict(tensor)

// ✅ 用 Worker 后台推理
const worker = new Worker('inference-worker.js')
worker.postMessage({ tensor: tensorData })
worker.onmessage = (e) => {
  const predictions = e.data
}

陷阱 3：输入格式错误

每个模型期望的输入格式不同（RGB vs BGR、归一化 etc）。查看文档仔细对照。

9. 与后端配合的最佳实践

虽然前端可以推理，但不是所有场景都适合：

典型架构：

用户上传图片 → 前端本地预处理
            → 前端跑小模型快速分类
            → 如果特殊情况→ 发送到后端做精准推理
            → 结果返回用户

10. 推荐工具与资源

• TensorFlow.js 官网：https://www.tensorflow.org/js
• 预训练模型库：https://github.com/tensorflow/tfjs-models
• 性能优化指南：TensorFlow Lite for Web（量化模型）
• Colab：快速原型开发

相关阅读

• 浏览器端图像识别实战
• Web Worker 性能优化
• 前端隐私保护最佳实践