JS `Machine Learning` `ONNX Runtime Web` `WebGPU` `Inference` `Optimization`

咳咳,各位观众老爷们,晚上好!今天咱们来聊聊一个既高大上又接地气的话题:用 JavaScript 搞机器学习,然后让它在浏览器里飞起来!具体来说,就是借助 ONNX Runtime Web + WebGPU 这对黄金搭档,把模型推理性能榨干最后一滴油。

一、机器学习:从 Python 到 JavaScript,换个姿势玩耍

说到机器学习,大家脑海里浮现的肯定是 Python,TensorFlow,PyTorch 这些关键词。没错,它们在服务器端叱咤风云,但在前端世界,JavaScript 才是王者。为什么要在浏览器里搞机器学习呢?原因很简单:

  • 离线推理: 用户不用再依赖服务器,网络不好也能愉快地使用你的应用。
  • 隐私保护: 数据处理在本地进行,不用担心敏感信息泄露。
  • 降低延迟: 减少网络传输,响应速度更快,用户体验更好。

当然,JavaScript 在机器学习领域也面临着挑战:

  • 性能瓶颈: JavaScript 本身不是为高性能计算设计的。
  • 生态不够完善: 相比 Python,JavaScript 机器学习库还比较少。

不过,这些问题都在逐渐被解决。今天咱们的主角 ONNX Runtime Web + WebGPU 就是来拯救性能的!

二、ONNX:模型界的“通用语”

ONNX (Open Neural Network Exchange) 是一个开放的模型格式标准,旨在让不同的机器学习框架能够互操作。你可以把 ONNX 理解为模型界的“通用语”,让 TensorFlow、PyTorch 等框架训练出来的模型,都能在不同的平台上运行。

有了 ONNX,我们就可以把 Python 训练好的模型导出成 ONNX 格式,然后在 JavaScript 中加载并运行。

三、ONNX Runtime Web:模型推理的“浏览器引擎”

ONNX Runtime Web 是 ONNX Runtime 的 Web 版本,它是一个高性能的推理引擎,专门用于在浏览器中运行 ONNX 模型。它可以利用 WebAssembly (WASM) 或 WebGPU 来加速推理过程。

  • WebAssembly (WASM): 一种二进制指令格式,可以在浏览器中以接近原生代码的速度运行。
  • WebGPU: 一种新的 Web API,允许 JavaScript 直接访问 GPU 资源,实现更强大的并行计算能力。

ONNX Runtime Web 的优势:

  • 跨平台: 可以在各种浏览器和操作系统上运行。
  • 高性能: 利用 WASM 或 WebGPU 加速推理。
  • 易于集成: 提供了简单的 JavaScript API。

四、WebGPU:让你的模型在 GPU 上“飙车”

WebGPU 是 W3C 组织推出的一个 Web API,旨在让 Web 应用能够利用 GPU 的强大计算能力。 相比之前的 WebGL,WebGPU 提供了更底层的接口,可以更好地控制 GPU 资源,从而实现更高的性能。

WebGPU 的优势:

  • 高性能: 充分利用 GPU 的并行计算能力。
  • 低开销: 减少 CPU 的负担。
  • 现代化: 采用了更现代化的 API 设计。

五、实战演练:用 ONNX Runtime Web + WebGPU 推理图像分类模型

接下来,咱们通过一个实际的例子,演示如何使用 ONNX Runtime Web + WebGPU 来推理一个图像分类模型。

1. 准备 ONNX 模型

首先,我们需要一个 ONNX 格式的图像分类模型。你可以自己训练一个,或者从 ONNX Model Zoo 下载一个现成的。这里我们以 SqueezeNet 为例。

2. 安装 ONNX Runtime Web

可以通过 npm 安装 ONNX Runtime Web:

npm install onnxruntime-web

3. HTML 结构

创建一个 HTML 文件,包含一个 canvas 用于显示图像,一个按钮用于选择图像,以及一个用于显示推理结果的区域。

<!DOCTYPE html>
<html>
<head>
  <title>ONNX Runtime Web Demo</title>
  <style>
    #imageCanvas {
      width: 224px;
      height: 224px;
      border: 1px solid black;
    }
  </style>
</head>
<body>
  <h1>ONNX Runtime Web Demo</h1>
  <input type="file" id="imageInput">
  <canvas id="imageCanvas"></canvas>
  <div id="result"></div>
  <script src="index.js"></script>
</body>
</html>

4. JavaScript 代码

index.js 文件中,编写 JavaScript 代码来加载模型、预处理图像、执行推理,并显示结果。

// 导入 ONNX Runtime Web
import * as ort from 'onnxruntime-web';

// 模型 URL
const modelUrl = 'squeezenet1.1.onnx'; // 替换为你的模型 URL

// 获取 HTML 元素
const imageInput = document.getElementById('imageInput');
const imageCanvas = document.getElementById('imageCanvas');
const resultDiv = document.getElementById('result');
const canvasContext = imageCanvas.getContext('2d');

// ONNX Runtime 会话
let session = null;

// 初始化 ONNX Runtime 会话
async function init() {
  try {
    // 创建 WebGPU 会话选项
    const options = {
      executionProviders: ['webgpu'], // 指定使用 WebGPU
      graphOptimizationLevel: 'all', // 优化图结构
    };

    // 创建 ONNX Runtime 会话
    session = await ort.InferenceSession.create(modelUrl, options);
    console.log('ONNX Runtime session created.');
  } catch (e) {
    console.error('Failed to create ONNX Runtime session:', e);
    resultDiv.innerText = 'Failed to load model. Check console for details.';
  }
}

// 图片上传事件处理
imageInput.addEventListener('change', async (event) => {
  const file = event.target.files[0];
  if (!file) return;

  const image = new Image();
  image.onload = async () => {
    // 在 canvas 上绘制图像
    canvasContext.drawImage(image, 0, 0, 224, 224);

    // 预处理图像
    const imageData = canvasContext.getImageData(0, 0, 224, 224);
    const inputTensor = preprocessImage(imageData);

    // 执行推理
    const outputTensor = await runInference(inputTensor);

    // 后处理结果
    const results = postprocessResults(outputTensor);

    // 显示结果
    displayResults(results);
  };
  image.src = URL.createObjectURL(file);
});

// 预处理图像
function preprocessImage(imageData) {
  const { data, width, height } = imageData;
  const red = [];
  const green = [];
  const blue = [];

  for (let i = 0; i < data.length; i += 4) {
    red.push(data[i] / 255.0);
    green.push(data[i + 1] / 255.0);
    blue.push(data[i + 2] / 255.0);
  }

  const input = [...red, ...green, ...blue];
  const tensor = new ort.Tensor('float32', input, [1, 3, width, height]);
  return tensor;
}

// 执行推理
async function runInference(inputTensor) {
  try {
    if (!session) {
      resultDiv.innerText = 'Model not loaded yet.';
      return null;
    }

    const feeds = { data: inputTensor };
    const results = await session.run(feeds);
    return results.softmaxout; // 假设输出是 'softmaxout'
  } catch (e) {
    console.error('Inference failed:', e);
    resultDiv.innerText = 'Inference failed. Check console for details.';
    return null;
  }
}

// 后处理结果
function postprocessResults(outputTensor) {
  if (!outputTensor) return [];

  const output = outputTensor.data;
  const topK = 5;
  const sorted = [...output]
    .map((value, index) => ({ value, index }))
    .sort((a, b) => b.value - a.value)
    .slice(0, topK);

  return sorted;
}

// 显示结果
function displayResults(results) {
  if (!results || results.length === 0) {
    resultDiv.innerText = 'No results to display.';
    return;
  }

  const labels = [
    "tench, Tinca tinca",
    "goldfish, Carassius auratus",
    "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias",
    "tiger shark, Galeocerdo cuvieri",
    "hammerhead, hammerhead shark",
    "electric ray, crampfish, numbfish, torpedo",
    "stingray",
    "cock",
    "hen",
    "ostrich, Struthio camelus",
    "brambling, Fringilla montifringilla",
    "goldfinch, Carduelis carduelis",
    "house finch, linnet, Carpodacus mexicanus",
    "junco, snowbird",
    "indigo bunting, indigo finch, indigo bird, Passerina cyanea",
    "robin, American robin, Turdus migratorius",
    "bulbul",
    "jay",
    "magpie",
    "chickadee",
    "water ouzel, dipper",
    "kite",
    "bald eagle, American eagle, Haliaeetus leucocephalus",
    "vulture",
    "great grey owl, great gray owl, Strix nebulosa",
    "European fire salamander, Salamandra salamandra",
    "common newt, Triturus vulgaris",
    "eft",
    "spotted salamander, Ambystoma maculatum",
    "axolotl, mud puppy, Ambystoma mexicanum",
    "bullfrog, Rana catesbeiana",
    "tree frog, tree-frog",
    "tailed frog, bell toad, ribbed toad, Ascaphus trui",
    "loggerhead, loggerhead turtle, Caretta caretta",
    "leatherback turtle, leatherback, leathery turtle, Dermochelys coriacea",
    "mud turtle",
    "terrapin",
    "box turtle, box tortoise",
    "banded gecko",
    "common iguana, iguana, Iguana iguana",
    "American chameleon, anole, Anolis carolinensis",
    "whiptail, whiptail lizard",
    "agama",
    "frilled lizard, Chlamydosaurus kingi",
    "alligator lizard",
    "Gila monster, Heloderma suspectum",
    "green lizard, Lacerta viridis",
    "African chameleon, Chamaeleo chamaeleon",
    "Komodo dragon, Komodo lizard, dragon lizard, giant lizard, Varanus komodoensis",
    "African crocodile, Nile crocodile, Crocodylus niloticus",
    "American alligator, Alligator mississipiensis",
    "triceratops",
    "thunder snake, worm snake, Carphophis amoenus",
    "ringneck snake, ring-necked snake, Diadophis punctatus",
    "hognose snake, puff adder, sand viper",
    "green snake, grass snake",
    "king snake, kingsnake",
    "garter snake, grass snake",
    "water snake",
    "vine snake",
    "night snake, Hypsiglena torquata",
    "boa constrictor, Constrictor constrictor",
    "rock python, rock snake, Python sebae",
    "Indian cobra, Naja naja",
    "green mamba",
    "sea snake",
    "horned viper, cerastes",
    "diamondback, diamondback rattlesnake, Crotalus adamanteus",
    "sidewinder, horned rattlesnake, Crotalus cerastes",
    "trilobite",
    "harvestman, daddy longlegs, Phalangium opilio",
    "scorpion",
    "black and gold garden spider, Argiope aurantia",
    "barn spider, Araneus cavaticus",
    "garden spider, Aranea diademata",
    "black widow, Latrodectus mactans",
    "tarantula",
    "wolf spider, hunting spider",
    "tick",
    "centipede",
    "black grouse",
    "ptarmigan",
    "ruffed grouse, partridge, Bonasa umbellus",
    "prairie chicken, prairie grouse, prairie fowl",
    "peacock",
    "quail",
    "partridge",
    "African grey, African gray, Psittacus erithacus",
    "macaw",
    "sulphur-crested cockatoo, Kakatoe galerita, Cacatua galerita",
    "lorikeet",
    "coucal",
    "bee eater",
    "hornbill",
    "hummingbird",
    "jacamar",
    "toucan",
    "drake",
    "red-breasted merganser, Mergus serrator",
    "goose",
    "black swan, Cygnus atratus",
    "tusker",
    "echidna, spiny anteater, anteater",
    "platypus, duckbill, duckbilled platypus, duck-billed platypus, Ornithorhynchus anatinus",
    "wallaby, brush kangaroo",
    "koala, koala bear, kangaroo bear, native bear, Phascolarctos cinereus",
    "wombat",
    "jellyfish",
    "sea anemone, anemone",
    "coral reef",
    "sturgeon",
    "gar, garfish, garpike, billfish, Lepisosteus osseus",
    "lionfish",
    "puffer, pufferfish, blowfish, globefish",
    "abacus",
    "abaya",
    "academic gown, academic robe, judge's robe",
    "accordion, piano accordion, squeeze box",
    "acoustic guitar",
    "aircraft carrier, carrier, flattop, attack aircraft carrier",
    "airliner",
    "airship, dirigible",
    "altar",
    "ambulance",
    "amphibian, amphibious vehicle",
    "analog clock",
    "apiary, bee house",
    "apron",
    "trash can, garbage can, wastebin, ashcan, trash barrel, trash bin, ash bin, ash-bin, garbage pail, waste basket, wastepaper basket, trash sack, garbage bag",
    "assault rifle, assault gun",
    "backpack, back pack, knapsack, packsack, rucksack, haversack",
    "bakery, bakeshop, bakehouse",
    "balance beam, beam",
    "balloon",
    "ballpoint, ballpoint pen, ballpen, Biro",
    "Band Aid",
    "banjo",
    "bannister, banister, balustrade, balusters, handrail",
    "barbell",
    "barber chair",
    "barometer",
    "barrel, cask",
    "barrow, garden cart, lawn cart, wheelbarrow",
    "baseball",
    "basketball",
    "bassinet",
    "bassoon",
    "bathing cap, swimming cap",
    "bath towel",
    "bathtub, bathing tub, bath, tub",
    "beach wagon, station wagon, wagon, estate car, beach waggon, station waggon, waggon",
    "beacon, lighthouse, beacon light, pharos",
    "beaker",
    "bearskin, busby, shako",
    "beer bottle",
    "beer glass",
    "bell cote, bell cot",
    "bib",
    "bicycle-built-for-two, tandem bicycle, tandem",
    "billiard table, pool table, snooker table",
    "binder, ring-binder",
    "binoculars, field glasses, opera glasses",
    "birdhouse",
    "biscuit",
    "blackboard, chalkboard",
    "blanket, cover",
    "blimp, gasbag",
    "blouse",
    "boathouse",
    "bobsled, bob, bobsleigh",
    "bolo tie, bolo, bola tie, bola",
    "bonnet, poke bonnet",
    "bookcase",
    "bookshop, bookstore",
    "bottlecap",
    "bow",
    "bow tie, bow-tie, bowtie",
    "brass, memorial tablet, plaque",
    "brassiere, bra, bandeau",
    "breakwater, groin, groyne, mole, bulwark, seawall, jetty",
    "brick kiln, brickkiln, brick field",
    "bridge, span",
    "broccoli",
    "broom",
    "bucket, pail",
    "buckle",
    "bulletproof vest",
    "bullhorn, megaphone",
    "bullet train, bullet",
    "butcher shop, meat market, meat shop",
    "cab, hack, taxi, taxicab",
    "caldron, cauldron",
    "candle, taper, wax light",
    "cannon",
    "canoe",
    "can opener, tin opener",
    "cardigan",
    "car mirror",
    "carousel, carrousel, merry-go-round, roundabout, whirligig",
    "carpenter's kit, tool kit",
    "carton",
    "car wheel",
    "cash machine, cash dispenser, automated teller machine, automatic teller machine, ATM",
    "cassette",
    "cassette player",
    "castle",
    "catamaran",
    "CD player",
    "cello, violoncello",
    "cellular telephone, cellular phone, cellphone, cell, mobile phone",
    "chain",
    "chain saw, chainsaw",
    "chest",
    "chiffonier, commode",
    "chime, bell, gong",
    "china cabinet, china closet",
    "Christmas stocking",
    "church, church building",
    "cinema, movie theater, movie theatre, movie house, picture palace",
    "cleaver, meat cleaver, chopper",
    "cliff, drop, drop-off",
    "cloak",
    "clog, geta, patten, sabot",
    "coffeepot",
    "coil, spiral, volute, whorl, helix",
    "combination lock",
    "computer keyboard, keypad",
    "confectionery, confectionary, candy store",
    "container ship, containership, container vessel",
    "convertible",
    "corkscrew, bottle screw",
    "cornet, horn, trumpet, trump",
    "cowboy boot",
    "crane",
    "crash helmet",
    "crate",
    "crib, cot",
    "Crock Pot",
    "croquet ball",
    "crutch",
    "cuirass",
    "dam, dike, dyke",
    "desk",
    "desktop computer",
    "dial telephone, dial phone",
    "diaper, nappy, napkin",
    "digital clock",
    "digital watch",
    "dining table, board",
    "dishcloth, dish towel",
    "dishwasher, dish washer, dishwashing machine",
    "disk brake, disc brake",
    "dock, dockage, docking facility",
    "dog sled, dog sleigh, husky sled",
    "dome",
    "doormat, welcome mat",
    "drilling platform, offshore rig",
    "drum, membranophone, tympan",
    "drumstick",
    "dumbbell",
    "Dutch oven",
    "electric fan, blower",
    "electric guitar",
    "electric locomotive",
    "entertainment center",
    "envelope",
    "espresso maker",
    "face powder",
    "feather boa, boa",
    "file, file cabinet, filing cabinet",
    "fireboat",
    "fire engine, fire truck",
    "fire screen, fireguard",
    "flagpole, flagstaff",
    "flute, transverse flute",
    "folding chair",
    "football helmet",
    "forklift",
    "fountain",
    "fountain pen",
    "four-poster",
    "freight car",
    "French horn, horn",
    "frying pan, frypan, skillet",
    "fur coat",
    "garbage truck, dustcart",
    "gasmask, respirator, gas helmet",
    "gas pump, gasoline pump, petrol pump, island dispenser",
    "goblet",
    "go-kart, kart",
    "golf ball",
    "golfcart, golf cart",
    "gondola",
    "gong, tam-tam",
    "gown",
    "grand piano, grand",
    "greenhouse, nursery, glasshouse",
    "grille, radiator grille",
    "grocery store, grocery, food market, market",
    "guillotine",
    "hair slide",
    "hair spray",
    "half track",
    "hammer",
    "hamper",
    "hand blower, blow dryer, blow drier, hair dryer, hair drier",
    "hand-held computer, hand-held device",
    "handkerchief, hankie, hanky, kerchief",
    "hard disc, hard disk, fixed disk",
    "harmonica, mouth organ, harp, mouth harp",
    "harp",
    "harvester, reaper",
    "hatchet",
    "holster",
    "home theater, home theatre",
    "honeycomb",
    "hook, claw",
    "hoopskirt, crinoline",
    "horizontal bar, high bar",
    "horse cart, horse-drawn vehicle",
    "hourglass",
    "iPod",
    "iron, smoothing iron",
    "jack-o'-lantern",
    "jean, blue jean, denim",
    "jeep, landrover",
    "jersey, T-shirt, tee shirt",
    "jigsaw puzzle",
    "jinrikisha, ricksha, rickshaw",
    "joystick",
    "kimono",
    "knee pad",
    "knot",
    "lab coat, laboratory coat",
    "ladle",
    "lampshade, lamp shade",
    "laptop, laptop computer",
    "lawn mower, mower",
    "lens cap, lens cover",
    "letter opener, paper knife, paperknife",
    "library",
    "lifeboat",
    "lighter, light, igniter, ignitor",
    "limousine, limo",
    "liner, ocean liner",
    "lipstick, lip rouge",
    "lotion",
    "loudspeaker, speaker, speaker unit, loudspeaker system",
    "loupe, jeweler's loupe",
    "lumbermill, sawmill",
    "magnetic compass",
    "mailbag, postbag",
    "mailbox, letter box",
    "maillot",
    "manhole cover",
    "maraca",
    "marimba, xylophone",
    "mask",
    "matchstick",
    "maypole",
    "maze, labyrinth",
    "measuring cup",
    "medicine chest, medicine cabinet",
    "megalith, megalithic structure",
    "microphone, mike",
    "microwave, microwave oven",
    "military uniform",
    "milk can",
    "minibus",
    "miniskirt, mini",
    "minivan",
    "missile",
    "mitten",
    "mixing bowl",
    "mobile home, manufactured home",
    "Model T",
    "modem",
    "monastery",
    "monitor",
    "moped",
    "mortar",
    "mortarboard",
    "mosque",
    "mosquito net",
    "motor scooter, scooter",
    "mountain bike, all-terrain bike, off-roader",
    "mountain tent",
    "mouse, computer mouse",
    "mousetrap",
    "moving van",
    "muzzle",
    "nail",
    "neck brace",
    "necklace",
    "nipple",
    "notebook, notepad",
    "obelisk",
    "oboe, hautboy, hautboise",
    "ocarina, sweet potato",
    "odometer, hodometer, mileometer, milometer",
    "oil filter",
    "organ, pipe organ",
    "oscilloscope, scope, cathode-ray oscilloscope, CRO",
    "overskirt",
    "oxcart",
    "oxygen mask",
    "packet",
    "paddle, boat paddle",
    "paddlewheel, paddle wheel",
    "padlock",
    "paintbrush",
    "pajama, pyjama, pj's, jammies",
    "palace",
    "panpipe, pandean pipe, syrinx",
    "paper towel",
    "parachute, chute",
    "parallel bars, bars",
    "park bench",
    "parking meter",
    "passenger car, coach, carriage",
    "patio, terrace",
    "pay-phone, pay-station",
    "pedestal, plinth, footstall",
    "pencil box, pencil case",
    "pencil sharpener",
    "perfume, essence",
    "Petri dish",
    "photocopier",
    "piano, pianoforte, forte-piano",
    "pick, plectrum, plectron",
    "pickelhaube",
    "picket fence, paling",
    "pickup, pickup truck",
    "pier",
    "piggy bank, penny bank",
    "pill bottle",
    "pillow",
    "ping-pong ball",
    "pinwheel",
    "pirate, pirate ship",
    "pitcher, ewer",
    "plane, carpenter's plane, woodworking plane",
    "planetarium",
    "plastic bag",
    "plate rack",
    "plow, plough",
    "plunger, plumber's helper",
    "Polaroid camera, Polaroid Land camera",
    "pole",
    "police van, police wagon, paddy wagon, patrol wagon, wagon, black Maria",
    "poncho",
    "pool table, billiard table, snooker table",
    "pop bottle, soda bottle",
    "pot, flowerpot",
    "potter's wheel",
    "power drill",
    "prayer rug, prayer mat",
    "printer",
    "prison, jail",
    "projectile, missile",
    "projector",
    "puck, hockey puck",
    "punching bag, punch bag, punching ball, punchball",
    "purse",
    "quill, quill pen",
    "quilt, comforter, comfort, puff",
    "racer, race car, racing car",
    "radiator",
    "radio, wireless",
    "radio telescope, radio reflector",
    "rain barrel",
    "ram, Tupa",
    "range hood, exhaust hood",
    "rapeseed",
    "reamer",
    "remote control, remote",
    "restaurant, eating house, eating place, eatery",
    "revolver, six-gun, six-shooter",
    "rifle",
    "rocking chair, rocker",
    "rotisserie",
    "rubber eraser, rubber, pencil eraser",
    "rugby ball",
    "rule, ruler",
    "running shoe",
    "safe",
    "safety goggles",
    "sail, canvas",
    "sailboat, sailing boat",
    "sandal",
    "sandwich",
    "sarong",
    "sax, saxophone",
    "scabbard",
    "scale, weighing machine",
    "school bus",
    "schooner",
    "scoreboard",
    "screen, CRT screen",
    "screw",
    "screwdriver",
    "seat belt, seatbelt",
    "sewing machine",
    "shield, buckler",
    "shoe shop, shoe-shop, shoe store",
    "shoji",
    "shopping basket",
    "shopping cart",
    "shovel",
    "shower cap",
    "shower curtain",
    "ski",
    "ski mask",
    "sleeping bag",
    "slide rule, slipstick",
    "sliding door",
    "slot machine",
    "snorkel",
    "snowmobile",
    "snowplow, snowplough",
    "soap dispenser",
    "soccer ball",
    "sock",
    "solar dish, solar collector, solar furnace",
    "sombrero",
    "soup bowl",
    "space bar",
    "space heater",
    "space shuttle",
    "spatula",
    "speedboat",
    "spider web, spider's web",
    "spindle",
    "sports car, sport car",
    "spotlight, spot",
    "stage",
    "steam locomotive",
    "steel arch bridge",
    "steel drum",
    "stethoscope",
    "stole",
    "stone wall",
    "stopwatch, stop watch",
    "stove",
    "strainer",
    "streetcar, tram, tramcar, trolley, trolley car",
    "stretcher",
    "studio couch, day bed",
    "stupa, tope",
    "submarine, pigboat, sub, U-boat",
    "suit, suit of clothes",
    "sundial",
    "sunglass",
    "sunglasses, dark glasses, shades",
    "sunscreen, sunblock, sun blocker",
    "suspension bridge",
    "swab, cotton swab",
    "sweatshirt",
    "swimming trunks, bathing trunks",
    "swing",
    "switch, electric switch, electrical switch",
    "syringe",
    "table lamp",
    "tank, army tank, armored combat vehicle, armoured combat vehicle",
    "tape player",
    "teapot",
    "teddy, teddy bear",
    "television, television system",
    "tennis ball",
    "thatch, thatched roof",
    "theater curtain, theatre curtain",
    "thimble",
    "thresher, thrasher, threshing machine",
    "throne",
    "tile roof",
    "toaster",
    "tobacco shop, tobacconist shop, tobacconist",
    "toilet seat",
    "torch",
    "totem pole",
    "tow truck, tow car, wrecker",
    "toyshop",
    "tractor",
    "trailer truck, tractor trailer, trucking rig, rig, articulated lorry, semi",
    "tray",
    "trench coat",
    "tricycle, trike, velocipede",
    "trimaran",
    "tripod",
    "triumphal arch",
    "trolleybus, trolley coach, trackless trolley",
    "trombone",
    "tub, vat",
    "turnstile",
    "typewriter keyboard",
    "umbrella",
    "unicycle, monocycle",
    "upright, upright piano",
    "vacuum cleaner, Hoover",
    "vase",
    "vault",
    "velvet",
    "vending machine",
    "vestment",
    "viaduct",
    "violin, fiddle",
    "virtual reality device, head-mounted display",
    "volleyball",
    "waffle iron",
    "wall clock",
    "wallet, billfold, notecase, pocketbook",
    "wardrobe, closet, press",
    "warplane, military plane",
    "washbasin, handbasin, washbowl, lavabo, wash-hand basin",
    "washer, automatic washer, washing machine",
    "water bottle",
    "water jug",
    "water tower",
    "whiskey jug",
    "whistle",
    "wig",
    "window screen",
    "wine bottle",
    "wing",
    "wok",
    "wooden spoon",
    "wool, woolen, woollen",
    "worm fence, snake fence, snake-rail fence, Virginia fence",
    "wreck",
    "yawl",
    "yurt",
    "zebra"
  ];
  let resultText = '';
  results.forEach((result) => {
    const label = labels[result.index];
    const probability = result.value.toFixed(4);
    resultText += `${label}: ${probability}<br>`;
  });
  resultDiv.innerHTML = resultText;
}

// 初始化
init();

5. 运行代码

在浏览器中打开 HTML 文件,选择一张图片,就可以看到推理结果了。

六、性能优化:榨干 GPU 的最后一滴油

虽然 WebGPU 已经提供了很高的性能,但我们还可以通过一些技巧来进一步优化:

  • 模型量化: 将模型的权重从 float32 转换为 int8 或 uint8,可以减少模型的大小和计算量。
  • 图优化: ONNX Runtime Web 提供了图优化功能,可以合并一些计算节点,减少推理步骤。
  • 预热: 在正式推理之前,先运行几次推理,让 GPU 预热,可以提高后续推理的性能。
  • 线程池: 使用 Web Workers 创建线程池,将推理任务分配到不同的线程,可以利用多核 CPU 的优势。
  • 减少内存分配: 尽量避免在推理过程中频繁分配内存,可以使用对象池等技术来减少内存分配的开销。

七、表格总结

技术/工具 优势 劣势
ONNX 模型格式标准化,方便不同框架之间的模型转换。 需要进行模型转换,可能会引入一些误差。
ONNX Runtime Web 可以在浏览器中运行 ONNX 模型,支持 WASM 和 WebGPU 加速。 相比原生代码,性能仍然有差距。
WebGPU 充分利用 GPU 的并行计算能力,提供更高的性能。 兼容性问题,需要浏览器支持。API 相对底层,需要一定的学习成本。
模型量化 减少模型大小和计算量,提高推理速度。 可能会降低模型的精度。
图优化 合并计算节点,减少推理步骤。 可能会增加模型的复杂性。
线程池 利用多核 CPU 的优势,提高推理速度。 需要管理线程的创建和销毁,可能会增加代码的复杂性。

八、未来展望

随着 WebAssembly 和 WebGPU 的不断发展,JavaScript 在机器学习领域的潜力将越来越大。未来,我们可以期待:

  • 更强大的 Web API: 更多的 Web API 将会提供更底层的硬件访问能力,让 JavaScript 能够

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注