DEV Community

Alex Spinov
Alex Spinov

Posted on

TensorFlow.js Has a Free API That Runs Machine Learning Models in the Browser

TensorFlow.js brings Google's ML framework to JavaScript. Train models in Node.js, run inference in the browser — no Python needed.

Pre-Trained Models: Zero Training

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

// Image classification
const model = await mobilenet.load();
const img = document.getElementById("product-image");
const predictions = await model.classify(img);
console.log(predictions);
// [{ className: 'laptop', probability: 0.92 }]
Enter fullscreen mode Exit fullscreen mode

Object Detection

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

const model = await cocoSsd.load();
const predictions = await model.detect(img);
predictions.forEach(p => {
  console.log(`${p.class}: ${(p.score * 100).toFixed(1)}%`);
  console.log(`Box: [${p.bbox.join(", ")}]`);
});
Enter fullscreen mode Exit fullscreen mode

Text Toxicity Detection

import * as toxicity from "@tensorflow-models/toxicity";

const model = await toxicity.load(0.9);
const predictions = await model.classify(["This product is terrible!"]);
predictions.forEach(p => {
  console.log(`${p.label}: ${p.results[0].match}`);
});
Enter fullscreen mode Exit fullscreen mode

Custom Model: Train in Browser

// Define model
const model = tf.sequential({
  layers: [
    tf.layers.dense({ inputShape: [4], units: 16, activation: "relu" }),
    tf.layers.dense({ units: 8, activation: "relu" }),
    tf.layers.dense({ units: 1, activation: "sigmoid" }),
  ],
});

model.compile({ optimizer: "adam", loss: "binaryCrossentropy", metrics: ["accuracy"] });

// Train
const xs = tf.tensor2d(trainingData);
const ys = tf.tensor2d(labels);

await model.fit(xs, ys, {
  epochs: 100,
  batchSize: 32,
  validationSplit: 0.2,
  callbacks: {
    onEpochEnd: (epoch, logs) => console.log(`Epoch ${epoch}: loss=${logs.loss.toFixed(4)}`),
  },
});

// Predict
const prediction = model.predict(tf.tensor2d([[1, 2, 3, 4]]));
console.log(prediction.dataSync()); // [0.87]
Enter fullscreen mode Exit fullscreen mode

Tensor Operations

const a = tf.tensor2d([[1, 2], [3, 4]]);
const b = tf.tensor2d([[5, 6], [7, 8]]);

const sum = a.add(b);
const product = a.matMul(b);
const mean = a.mean();

console.log(sum.arraySync());     // [[6, 8], [10, 12]]
console.log(product.arraySync()); // [[19, 22], [43, 50]]
Enter fullscreen mode Exit fullscreen mode

Save and Load Models

// Save to localStorage
await model.save("localstorage://my-model");

// Save to file
await model.save("downloads://my-model");

// Load
const loaded = await tf.loadLayersModel("localstorage://my-model");
Enter fullscreen mode Exit fullscreen mode

ML-powered data analysis? My Apify tools + TensorFlow.js = intelligent scraping.

Custom ML solution? Email spinov001@gmail.com

Top comments (0)