TensorFlow.js for JavaScript Developers – Machine Learning in Browser

// Option 1: CDN — simplest for prototypes and demos
// <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.22.0"></script>

// Option 2: npm — for production bundlers (Next.js, Vite, Webpack)
// npm install @tensorflow/tfjs @tensorflow/tfjs-backend-webgl
import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-backend-webgl'; // Explicit backend import

// Verify installation, backend, and GPU availability
async function initTF() {
  await tf.ready();
  
  const backend = tf.getBackend();
  const memInfo = tf.memory();
  
  console.log(`TensorFlow.js v${tf.version.tfjs}`);
  console.log(`Backend: ${backend}`);
  console.log(`GPU Tensors: ${memInfo.numTensors}`);
  console.log(`GPU Memory: ${(memInfo.numBytes / 1e6).toFixed(1)}MB`);
  
  if (backend === 'cpu') {
    console.warn(
      'WARNING: Running on CPU backend. GPU acceleration is not available. ' +
      'Performance will be 10-50x slower than WebGL/WebGPU.'
    );
  }
  
  // Quick sanity test — verify tensor operations work
  const test = tf.tensor([1, 2, 3, 4]);
  console.log('Sanity check:', test.dataSync()); // [1, 2, 3, 4]
  test.dispose();
  
  return backend;
}

initTF();

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

// Load from URL (CDN-hosted model)
const MODEL_URL = 'https://storage.googleapis.com/your-bucket/models/classifier_v2/model.json';

async function loadModel() {
  await tf.ready();
  console.log(`Backend: ${tf.getBackend()}`);
  
  const startLoad = performance.now();
  
  const model = await tf.loadLayersModel(MODEL_URL, {
    onProgress: (fraction) => {
      // Update a loading bar in the UI
      console.log(`Loading: ${(fraction * 100).toFixed(1)}%`);
    }
  });
  
  const loadTime = performance.now() - startLoad;
  console.log(`Model loaded in ${loadTime.toFixed(0)}ms`);
  console.log(`Input shape: ${JSON.stringify(model.inputs[0].shape)}`);
  console.log(`Output shape: ${JSON.stringify(model.outputs[0].shape)}`);
  
  // Warm up — first prediction compiles GPU shaders
  const startWarmup = performance.now();
  const dummyInput = tf.zeros(model.inputs[0].shape.map(d => d || 1));
  const warmupOutput = model.predict(dummyInput);
  await warmupOutput.data(); // Force GPU sync — shaders compile here
  const warmupTime = performance.now() - startWarmup;
  
  tf.dispose([dummyInput, warmupOutput]);
  console.log(`Warm-up inference: ${warmupTime.toFixed(0)}ms (includes shader compilation)`);
  
  return model;
}

// Load from IndexedDB (cached model for offline and repeat visits)
async function loadCachedModel(modelId) {
  try {
    const model = await tf.loadLayersModel(`indexeddb://${modelId}`);
    console.log(`Loaded cached model: ${modelId}`);
    return model;
  } catch (err) {
    console.log(`No cached model found for ${modelId}, loading from network`);
    return null;
  }
}

// Save to IndexedDB after first network load
async function cacheModel(model, modelId) {
  await model.save(`indexeddb://${modelId}`);
  console.log(`Model cached as: ${modelId}`);
}

// Full loading strategy with cache-first pattern
async function loadModelWithCache(modelId, networkUrl) {
  // Try cache first
  let model = await loadCachedModel(modelId);
  
  if (!model) {
    // Cache miss — load from network
    model = await loadModel(networkUrl);
    await cacheModel(model, modelId);
  }
  
  // Warm up regardless of source
  const dummy = tf.zeros(model.inputs[0].shape.map(d => d || 1));
  const warm = model.predict(dummy);
  await warm.data();
  tf.dispose([dummy, warm]);
  
  return model;
}

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

// Image classification pipeline — single image
async function classifyImage(model, imageElement, labels) {
  // Preprocess: resize, normalize, add batch dimension
  // CRITICAL: normalization must match the Python training pipeline
  const inputTensor = tf.tidy(() => {
    return tf.browser.fromPixels(imageElement)  // [H, W, 3] uint8
      .resizeBilinear([224, 224])               // Match model input shape
      .toFloat()                                 // Cast to float32
      .div(127.5)                                // Scale to [0, 2]
      .sub(1.0)                                  // Shift to [-1, 1] (MobileNet convention)
      .expandDims(0);                            // Add batch dim: [1, 224, 224, 3]
  });

  // Run inference
  const predictions = model.predict(inputTensor);
  const probabilities = await predictions.data(); // GPU → CPU transfer

  // Cleanup — prevent memory leaks
  tf.dispose([inputTensor, predictions]);

  // Map to class labels and sort by confidence
  const results = Array.from(probabilities)
    .map((prob, i) => ({ label: labels[i], confidence: prob }))
    .sort((a, b) => b.confidence - a.confidence);

  return {
    topPrediction: results[0],
    allPredictions: results
  };
}

// Real-time video classification at target FPS
async function classifyVideoStream(model, videoElement, labels, targetFPS = 30) {
  const frameInterval = 1000 / targetFPS;
  let lastFrameTime = 0;
  let isProcessing = false;
  
  async function processFrame(timestamp) {
    // Skip frame if previous inference is still running
    if (isProcessing || timestamp - lastFrameTime < frameInterval) {
      requestAnimationFrame(processFrame);
      return;
    }
    
    isProcessing = true;
    lastFrameTime = timestamp;
    
    // All tensor ops wrapped in tidy for automatic cleanup
    const outputTensor = tf.tidy(() => {
      const frame = tf.browser.fromPixels(videoElement);
      const resized = tf.image.resizeBilinear(frame, [224, 224]);
      const normalized = resized.toFloat().div(127.5).sub(1.0);
      const batched = normalized.expandDims(0);
      return model.predict(batched);
    });

    const result = await outputTensor.data();
    outputTensor.dispose();
    
    // Use result — update UI, trigger action, etc.
    const topIndex = result.indexOf(Math.max(...result));
    console.log(`${labels[topIndex]}: ${(result[topIndex] * 100).toFixed(1)}%`);
    
    isProcessing = false;
    requestAnimationFrame(processFrame);
  }
  
  requestAnimationFrame(processFrame);
}

// Example: classify a file upload
async function handleFileUpload(model, file) {
  const img = new Image();
  img.src = URL.createObjectURL(file);
  await img.decode(); // Wait for image to load completely
  
  const labels = ['cat', 'dog', 'bird', 'fish', 'horse'];
  const result = await classifyImage(model, img, labels);
  
  URL.revokeObjectURL(img.src); // Clean up object URL
  return result;
}

# Install the converter
pip install tensorflowjs

# Convert Keras .h5 model with float16 quantization
tensorflowjs_converter \
  --input_format=keras \
  --output_format=tfjs_graph_model \
  --quantize_float16 \
  --weight_shard_size_bytes=4194304 \
  ./models/my_model.h5 \
  ./tfjs_models/my_model

# Convert SavedModel directory
tensorflowjs_converter \
  --input_format=tf_saved_model \
  --output_format=tfjs_graph_model \
  --signature_name=serving_default \
  --saved_model_tags=serve \
  --quantize_float16 \
  --weight_shard_size_bytes=4194304 \
  ./saved_model/ \
  ./tfjs_models/my_model

# Convert Keras .keras format (TF 2.16+)
tensorflowjs_converter \
  --input_format=keras \
  --output_format=tfjs_graph_model \
  --quantize_float16 \
  ./models/my_model.keras \
  ./tfjs_models/my_model

# Verify converted model output structure
ls -la ./tfjs_models/my_model/
# model.json              (graph topology + weight manifest)
# group1-shard1of3.bin    (weight data, ~4MB each)
# group1-shard2of3.bin
# group1-shard3of3.bin

# Check model size after conversion
du -sh ./tfjs_models/my_model/
# 23M    ./tfjs_models/my_model/   (from 46MB float32 original)

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

// Initialize with fallback chain: WebGPU → WebGL → CPU
async function initBestBackend() {
  const backends = ['webgpu', 'webgl', 'cpu'];
  
  for (const backend of backends) {
    try {
      // Feature detection for WebGPU
      if (backend === 'webgpu' && !('gpu' in navigator)) {
        console.log('WebGPU: not available in this browser');
        continue;
      }
      
      await tf.setBackend(backend);
      await tf.ready();
      console.log(`Backend initialized: ${backend}`);
      return backend;
    } catch (err) {
      console.warn(`${backend} backend failed: ${err.message}`);
    }
  }
  
  throw new Error('No TensorFlow.js backend available');
}

// Benchmark to compare backends on the actual device
async function benchmarkBackend(iterations = 10) {
  const backend = tf.getBackend();
  const a = tf.randomNormal([1024, 1024]);
  const b = tf.randomNormal([1024, 1024]);
  
  // Warm up — first run includes shader compilation
  const warmup = tf.matMul(a, b);
  await warmup.data();
  warmup.dispose();
  
  // Timed runs
  const times = [];
  for (let i = 0; i < iterations; i++) {
    const start = performance.now();
    const c = tf.matMul(a, b);
    await c.data(); // Force GPU sync
    times.push(performance.now() - start);
    c.dispose();
  }
  
  tf.dispose([a, b]);
  
  const avg = times.reduce((s, t) => s + t, 0) / times.length;
  const min = Math.min(...times);
  const max = Math.max(...times);
  
  console.log(`Backend: ${backend}`);
  console.log(`1024x1024 matMul (${iterations} runs):`);
  console.log(`  Avg: ${avg.toFixed(1)}ms`);
  console.log(`  Min: ${min.toFixed(1)}ms`);
  console.log(`  Max: ${max.toFixed(1)}ms`);
  
  return { backend, avg, min, max };
}

// Usage
const activeBackend = await initBestBackend();
await benchmarkBackend();

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

// PATTERN 1: tf.tidy for synchronous automatic cleanup
function predictSafely(model, imageElement) {
  // All intermediate tensors created inside tf.tidy are disposed automatically
  // Only the returned tensor survives
  return tf.tidy(() => {
    const input = tf.browser.fromPixels(imageElement)
      .toFloat()          // intermediate tensor 1
      .div(255.0)          // intermediate tensor 2
      .expandDims(0);      // intermediate tensor 3
    return model.predict(input);  // only this survives
  });
}

// PATTERN 2: Manual disposal for async operations
async function predictAsync(model, imageElement) {
  let input = null;
  let output = null;
  
  try {
    input = tf.tidy(() => {
      return tf.browser.fromPixels(imageElement)
        .toFloat().div(255.0).expandDims(0);
    });
    
    output = model.predict(input);
    const result = await output.data(); // async — cannot use tf.tidy for this
    return Array.from(result);
  } finally {
    // Dispose in finally block — runs even if an error is thrown
    if (input) input.dispose();
    if (output) output.dispose();
  }
}

// ANTI-PATTERN: Memory leak — tensors never disposed
function predictLeaky(model, imageElement) {
  // BAD: three intermediate tensors leak on every call
  const pixels = tf.browser.fromPixels(imageElement); // leaked
  const floats = pixels.toFloat();                     // leaked
  const normalized = floats.div(255.0);                // leaked
  const batched = normalized.expandDims(0);            // leaked
  const output = model.predict(batched);               // leaked
  return output.data();
  // Nothing is ever disposed — GPU memory grows until crash
}

// Memory monitoring — use in development to detect leaks
function assertNoLeaks(label, fn) {
  const before = tf.memory().numTensors;
  fn();
  const after = tf.memory().numTensors;
  if (after > before + 1) { // +1 for the returned tensor
    console.error(
      `[LEAK] ${label}: ${after - before} tensors created, ` +
      `expected at most 1. Before: ${before}, After: ${after}`
    );
  }
}

// Full lifecycle monitoring
function logMemory(label) {
  const info = tf.memory();
  console.log(
    `[${label}] Tensors: ${info.numTensors} | ` +
    `Bytes: ${(info.numBytes / 1e6).toFixed(1)}MB | ` +
    `Unreliable: ${info.unreliable}`
  );
}

// Cleanup when a model is no longer needed
function disposeModel(model) {
  model.dispose(); // Frees all weight tensors and GPU resources
  console.log('Model disposed. Remaining tensors:', tf.memory().numTensors);
}

// components/ImageClassifier.jsx
'use client';

import { useState, useEffect, useRef, useCallback } from 'react';

export default function ImageClassifier() {
  const [loading, setLoading] = useState(true);
  const [result, setResult] = useState(null);
  const [error, setError] = useState(null);
  const modelRef = useRef(null);
  const tfRef = useRef(null);

  useEffect(() => {
    let cancelled = false;

    async function init() {
      try {
        // Dynamic import of TensorFlow.js — only in browser
        const tf = await import('@tensorflow/tfjs');
        await import('@tensorflow/tfjs-backend-webgl');
        tfRef.current = tf;

        await tf.ready();
        console.log(`Backend: ${tf.getBackend()}`);

        const model = await tf.loadLayersModel('/models/classifier/model.json');

        // Warm up with dummy prediction
        const inputShape = model.inputs[0].shape.map(d => d || 1);
        const dummy = tf.zeros(inputShape);
        const warm = model.predict(dummy);
        await warm.data();
        tf.dispose([dummy, warm]);

        if (!cancelled) {
          modelRef.current = model;
          setLoading(false);
        } else {
          model.dispose(); // Component unmounted during loading
        }
      } catch (err) {
        if (!cancelled) {
          setError(err.message);
          setLoading(false);
        }
      }
    }

    init();

    // Cleanup on unmount — prevents GPU memory leak on route change
    return () => {
      cancelled = true;
      if (modelRef.current) {
        modelRef.current.dispose();
        modelRef.current = null;
        console.log('Model disposed on component unmount');
      }
    };
  }, []);

  const handlePredict = useCallback(async (imageElement) => {
    const model = modelRef.current;
    const tf = tfRef.current;
    if (!model || !tf) return null;

    const prediction = tf.tidy(() => {
      const input = tf.browser.fromPixels(imageElement)
        .resizeBilinear([224, 224])
        .toFloat().div(127.5).sub(1.0)
        .expandDims(0);
      return model.predict(input);
    });

    const data = await prediction.data();
    prediction.dispose();

    const labels = ['cat', 'dog', 'bird', 'fish', 'horse'];
    const topIndex = Array.from(data).indexOf(Math.max(...data));
    const newResult = { label: labels[topIndex], confidence: data[topIndex] };
    setResult(newResult);
    return newResult;
  }, []);

  if (error) return <p>ML Error: {error}</p>;
  if (loading) return <p>Loading ML model...</p>;
  return <div>Model ready. Result: {result?.label} ({(result?.confidence * 100)?.toFixed(1)}%)</div>;
}

// app/page.jsx — dynamic import prevents SSR
import dynamic from 'next/dynamic';

const ImageClassifier = dynamic(
  () => import('@/components/ImageClassifier'),
  {
    ssr: false,
    loading: () => <p>Initializing ML engine...</p>
  }
);

export default function Home() {
  return (
    <main>
      <h1>Browser ML Demo</h1>
      <ImageClassifier />
    </main>
  );
}

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

// Technique 1: Batch predictions for throughput
async function batchPredict(model, images) {
  // Stack multiple images into one batch tensor — one GPU dispatch
  const batchTensor = tf.tidy(() => {
    const tensors = images.map(img =>
      tf.browser.fromPixels(img)
        .resizeBilinear([224, 224])
        .toFloat().div(127.5).sub(1.0)
    );
    return tf.stack(tensors); // [batch, 224, 224, 3]
  });

  const predictions = model.predict(batchTensor);
  const results = await predictions.array();

  tf.dispose([batchTensor, predictions]);
  return results;
}

// Technique 2: Reduce input resolution for real-time speed
function preprocessAtResolution(imageElement, targetSize = 128) {
  return tf.tidy(() => {
    return tf.browser.fromPixels(imageElement)
      .resizeBilinear([targetSize, targetSize]) // 128x128 = 66% fewer pixels than 224x224
      .toFloat().div(127.5).sub(1.0)
      .expandDims(0);
  });
}

// Technique 3: Profile inference to find bottlenecks
async function profileInference(model, inputShape, runs = 20) {
  const input = tf.randomNormal(inputShape);

  // Warm up — exclude shader compilation from timing
  const warmup = model.predict(input);
  await warmup.data();
  warmup.dispose();

  // Timed runs
  const times = [];
  for (let i = 0; i < runs; i++) {
    const start = performance.now();
    const output = model.predict(input);
    await output.data(); // Force GPU sync
    times.push(performance.now() - start);
    output.dispose();
  }

  input.dispose();

  const avg = times.reduce((s, t) => s + t, 0) / times.length;
  const p95 = times.sort((a, b) => a - b)[Math.floor(times.length * 0.95)];
  const min = times[0];

  console.log(`Inference profile (${runs} runs, ${tf.getBackend()} backend):`);
  console.log(`  Average: ${avg.toFixed(1)}ms`);
  console.log(`  P95:     ${p95.toFixed(1)}ms`);
  console.log(`  Min:     ${min.toFixed(1)}ms`);
  console.log(`  Target:  ${avg < 33 ? '✓ 30 FPS achievable' : '✗ Too slow for 30 FPS'}`);

  return { avg, p95, min };
}

// Technique 4: Skip frames when inference cannot keep up
class AdaptiveInference {
  constructor(model, targetFPS = 30) {
    this.model = model;
    this.targetInterval = 1000 / targetFPS;
    this.isProcessing = false;
    this.lastTime = 0;
    this.droppedFrames = 0;
    this.processedFrames = 0;
  }

  async processFrame(imageElement, timestamp) {
    if (this.isProcessing) {
      this.droppedFrames++;
      return null; // Skip — previous frame still processing
    }

    if (timestamp - this.lastTime < this.targetInterval) {
      return null; // Skip — too soon since last frame
    }

    this.isProcessing = true;
    this.lastTime = timestamp;

    const output = tf.tidy(() => {
      const input = tf.browser.fromPixels(imageElement)
        .resizeBilinear([128, 128])
        .toFloat().div(127.5).sub(1.0)
        .expandDims(0);
      return this.model.predict(input);
    });

    const result = await output.data();
    output.dispose();

    this.processedFrames++;
    this.isProcessing = false;

    return result;
  }

  getStats() {
    const total = this.processedFrames + this.droppedFrames;
    return {
      processed: this.processedFrames,
      dropped: this.droppedFrames,
      dropRate: total > 0 ? (this.droppedFrames / total * 100).toFixed(1) + '%' : '0%'
    };
  }
}

// Usage
const profiler = await profileInference(model, [1, 224, 224, 3]);
const adaptive = new AdaptiveInference(model, 30);