TensorFlow.js 浏览器推理 #

┌─────────────────────────────────────────────────────────────┐
│                    浏览器推理优势                             │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐         │
│  │  零延迟     │  │  隐私保护   │  │  离线可用   │         │
│  │  无网络请求 │  │  数据本地   │  │  无需服务器 │         │
│  └─────────────┘  └─────────────┘  └─────────────┘         │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐         │
│  │  跨平台     │  │  实时交互   │  │  低成本     │         │
│  │  任何设备   │  │  即时反馈   │  │  无服务器费 │         │
│  └─────────────┘  └─────────────┘  └─────────────┘         │
│                                                             │
└─────────────────────────────────────────────────────────────┘

console.log('当前后端:', tf.getBackend());

await tf.setBackend('webgl');
await tf.setBackend('wasm');
await tf.setBackend('cpu');

await tf.ready();
console.log('后端就绪:', tf.getBackend());

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

await tf.setBackend('webgl');
await tf.ready();

console.log('WebGL 版本:', tf.ENV.get('WEBGL_VERSION'));
console.log('最大纹理尺寸:', tf.ENV.get('WEBGL_MAX_TEXTURE_SIZE'));

npm install @tensorflow/tfjs-backend-wasm

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

await tf.setBackend('wasm');
await tf.ready();

async function getBestBackend() {
  const backends = ['webgl', 'wasm', 'cpu'];
  
  for (const backend of backends) {
    try {
      await tf.setBackend(backend);
      await tf.ready();
      console.log(`使用后端: ${backend}`);
      return backend;
    } catch (e) {
      console.log(`后端 ${backend} 不可用`);
    }
  }
  
  return 'cpu';
}

let modelPromise = null;

function getModel() {
  if (!modelPromise) {
    modelPromise = tf.loadLayersModel('model.json');
  }
  return modelPromise;
}

async function predict(input) {
  const model = await getModel();
  return model.predict(input);
}

async function loadModelWithCache(url) {
  const cacheName = 'tfjs-models';
  const cache = await caches.open(cacheName);
  
  const cachedResponse = await cache.match(url);
  if (cachedResponse) {
    console.log('从缓存加载模型');
    return tf.loadLayersModel(url);
  }
  
  console.log('从网络加载模型');
  const model = await tf.loadLayersModel(url);
  
  await cache.add(url);
  
  return model;
}

async function loadModelProgressively(url, onProgress) {
  const response = await fetch(url);
  const contentLength = response.headers.get('content-length');
  const total = parseInt(contentLength, 10);
  let loaded = 0;
  
  const reader = response.body.getReader();
  const chunks = [];
  
  while (true) {
    const { done, value } = await reader.read();
    if (done) break;
    
    chunks.push(value);
    loaded += value.length;
    
    if (onProgress) {
      onProgress(loaded / total);
    }
  }
  
  const blob = new Blob(chunks);
  const modelUrl = URL.createObjectURL(blob);
  return tf.loadLayersModel(modelUrl);
}

function predictWithCleanup(model, input) {
  return tf.tidy(() => {
    const output = model.predict(input);
    return output.dataSync();
  });
}

async function batchPredict(model, inputs, batchSize = 32) {
  const results = [];
  
  for (let i = 0; i < inputs.length; i += batchSize) {
    const batch = inputs.slice(i, i + batchSize);
    const batchTensor = tf.stack(batch);
    
    const predictions = model.predict(batchTensor);
    results.push(...await predictions.array());
    
    batchTensor.dispose();
    predictions.dispose();
  }
  
  return results;
}

async function warmupModel(model, inputShape) {
  const dummyInput = tf.zeros(inputShape);
  await model.predict(dummyInput).data();
  dummyInput.dispose();
  console.log('模型预热完成');
}

async function predictAsync(model, input) {
  const output = await model.predictAsync(input);
  const result = await output.data();
  output.dispose();
  return result;
}

const workerCode = `
  importScripts('https://cdn.jsdelivr.net/npm/@tensorflow/tfjs');
  
  let model = null;
  
  self.onmessage = async (e) => {
    const { type, data } = e.data;
    
    if (type === 'init') {
      model = await tf.loadLayersModel(data.modelUrl);
      await model.predict(tf.zeros(data.inputShape)).data();
      self.postMessage({ type: 'ready' });
    }
    
    if (type === 'predict') {
      const input = tf.tensor(data.input);
      const output = model.predict(input);
      const result = await output.data();
      
      input.dispose();
      output.dispose();
      
      self.postMessage({ type: 'result', data: result });
    }
  };
`;

const blob = new Blob([workerCode], { type: 'application/javascript' });
const workerUrl = URL.createObjectURL(blob);
const worker = new Worker(workerUrl);

function initWorker(modelUrl, inputShape) {
  return new Promise((resolve) => {
    worker.onmessage = (e) => {
      if (e.data.type === 'ready') {
        resolve();
      }
    };
    
    worker.postMessage({
      type: 'init',
      data: { modelUrl, inputShape }
    });
  });
}

function predictInWorker(input) {
  return new Promise((resolve) => {
    worker.onmessage = (e) => {
      if (e.data.type === 'result') {
        resolve(e.data.data);
      }
    };
    
    worker.postMessage({
      type: 'predict',
      data: { input }
    });
  });
}

async function processVideoStream(model, video, callback) {
  const canvas = document.createElement('canvas');
  const ctx = canvas.getContext('2d');
  
  canvas.width = 224;
  canvas.height = 224;
  
  async function processFrame() {
    if (video.paused || video.ended) return;
    
    ctx.drawImage(video, 0, 0, 224, 224);
    const imageData = ctx.getImageData(0, 0, 224, 224);
    
    const tensor = tf.tidy(() => {
      return tf.browser.fromPixels(imageData)
        .toFloat()
        .div(255)
        .expandDims(0);
    });
    
    const prediction = model.predict(tensor);
    const result = await prediction.data();
    
    callback(result);
    
    tensor.dispose();
    prediction.dispose();
    
    requestAnimationFrame(processFrame);
  }
  
  processFrame();
}

async function setupCamera() {
  const video = document.getElementById('video');
  
  const stream = await navigator.mediaDevices.getUserMedia({
    video: { width: 640, height: 480 }
  });
  
  video.srcObject = stream;
  await video.play();
  
  return video;
}

async function realtimeInference() {
  const model = await tf.loadLayersModel('model.json');
  const video = await setupCamera();
  
  const fps = 30;
  const interval = 1000 / fps;
  let lastTime = 0;
  
  async function processFrame(timestamp) {
    if (timestamp - lastTime >= interval) {
      lastTime = timestamp;
      
      const tensor = tf.tidy(() => {
        return tf.browser.fromPixels(video)
          .resizeNearestNeighbor([224, 224])
          .toFloat()
          .div(255)
          .expandDims(0);
      });
      
      const prediction = model.predict(tensor);
      const result = await prediction.data();
      
      updateUI(result);
      
      tensor.dispose();
      prediction.dispose();
    }
    
    requestAnimationFrame(processFrame);
  }
  
  requestAnimationFrame(processFrame);
}

const model = await tf.loadLayersModel(
  'https://example.com/model_quantized.json'
);

tensorflowjs_converter \
  --input_format keras \
  --quantization_bytes 2 \
  model.h5 \
  ./quantized_model

async function loadShardedModel(baseUrl, shards) {
  const modelArtifacts = {
    modelTopology: null,
    weightSpecs: [],
    weightData: []
  };
  
  for (const shard of shards) {
    const response = await fetch(`${baseUrl}/${shard}`);
    const data = await response.json();
    
    if (shard === 'model.json') {
      modelArtifacts.modelTopology = data.modelTopology;
    } else {
      modelArtifacts.weightSpecs.push(...data.weightSpecs);
      modelArtifacts.weightData.push(...data.weightData);
    }
  }
  
  return tf.loadLayersModel(tf.io.fromMemory(modelArtifacts));
}

async function safeSetBackend(backend) {
  try {
    await tf.setBackend(backend);
    await tf.ready();
    return true;
  } catch (error) {
    console.error(`后端 ${backend} 初始化失败:`, error);
    return false;
  }
}

async function loadModelWithFallback(primaryUrl, fallbackUrl) {
  try {
    return await tf.loadLayersModel(primaryUrl);
  } catch (error) {
    console.error('主模型加载失败，尝试备用模型');
    return await tf.loadLayersModel(fallbackUrl);
  }
}

async function predictWithMemoryCheck(model, input) {
  const memBefore = tf.memory();
  
  try {
    const output = model.predict(input);
    const result = await output.data();
    output.dispose();
    return result;
  } catch (error) {
    if (error.message.includes('memory')) {
      console.error('内存不足，清理后重试');
      tf.disposeVariables();
      return predictWithMemoryCheck(model, input);
    }
    throw error;
  }
}

function monitorMemory() {
  const mem = tf.memory();
  console.log({
    张量数量: mem.numTensors,
    字节数: (mem.numBytes / 1024 / 1024).toFixed(2) + ' MB',
    不稳定张量: mem.numDataBuffers
  });
}

setInterval(monitorMemory, 1000);

async function measureInferenceTime(model, input, runs = 100) {
  const times = [];
  
  for (let i = 0; i < runs; i++) {
    const start = performance.now();
    const output = model.predict(input);
    await output.data();
    const end = performance.now();
    
    times.push(end - start);
    output.dispose();
  }
  
  const avg = times.reduce((a, b) => a + b) / times.length;
  const min = Math.min(...times);
  const max = Math.max(...times);
  
  console.log(`平均: ${avg.toFixed(2)}ms, 最小: ${min.toFixed(2)}ms, 最大: ${max.toFixed(2)}ms`);
}

class FPSMonitor {
  constructor(windowSize = 60) {
    this.frames = [];
    this.windowSize = windowSize;
    this.lastTime = performance.now();
  }
  
  tick() {
    const now = performance.now();
    const delta = now - this.lastTime;
    this.lastTime = now;
    
    this.frames.push(delta);
    if (this.frames.length > this.windowSize) {
      this.frames.shift();
    }
  }
  
  getFPS() {
    const avgDelta = this.frames.reduce((a, b) => a + b) / this.frames.length;
    return 1000 / avgDelta;
  }
}

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

class BrowserInference {
  constructor() {
    this.model = null;
    this.backend = null;
    this.isReady = false;
  }
  
  async init(modelUrl, preferredBackend = 'webgl') {
    await this.setBestBackend(preferredBackend);
    this.model = await tf.loadLayersModel(modelUrl);
    await this.warmup();
    this.isReady = true;
  }
  
  async setBestBackend(preferred) {
    const backends = [preferred, 'wasm', 'cpu'];
    
    for (const backend of backends) {
      try {
        await tf.setBackend(backend);
        await tf.ready();
        this.backend = backend;
        console.log(`使用后端: ${backend}`);
        return;
      } catch (e) {
        console.log(`后端 ${backend} 不可用`);
      }
    }
  }
  
  async warmup() {
    const inputShape = this.model.inputShape.slice(1);
    const dummyInput = tf.zeros([1, ...inputShape]);
    await this.model.predict(dummyInput).data();
    dummyInput.dispose();
  }
  
  async predict(input, options = {}) {
    if (!this.isReady) {
      throw new Error('模型未初始化');
    }
    
    return tf.tidy(() => {
      const tensor = tf.tensor(input).expandDims(0);
      const output = this.model.predict(tensor);
      return output.dataSync();
    });
  }
  
  getMemoryInfo() {
    return tf.memory();
  }
  
  dispose() {
    if (this.model) {
      this.model.dispose();
    }
  }
}

const inference = new BrowserInference();
await inference.init('model.json');
const result = await inference.predict([1, 2, 3, 4]);
console.log(result);