Web AI Model Testing: WebGPU, WebGL, and Headless Chrome

Great news! You’ve built a cool Web AI application
that runs machine learning models directly on a user’s device. It runs entirely
on the client-side web browser, without relying on the cloud. This on-device
design enhances user privacy, boosts performance, and reduces costs
significantly.

However, there’s a hurdle. Your
TensorFlow.js model can operate on
both CPUs (WebAssembly) and more powerful GPUs (through
WebGL and
WebGPU). The question is:
how can you consistently automate browser testing with the selected hardware?

Maintaining consistency is crucial for comparing machine learning model
performance over time as you iterate and improve them, prior to deployment for
real-world users to use on their device.

Setting up a consistent testing environment with GPUs can be harder than
expected. In this blog post, we’ll share the problems we faced and how we solved
them, so you can improve your application’s performance.

This isn’t just for Web AI developers! If you’re working on web gaming or
graphics, this post is valuable for you, too.

What’s in our automation toolbox

Here is what we are using:

• Environment: A Linux-based Google Colab
  notebook connected to an NVIDIA
  T4 or V100 GPU. You can use other cloud platforms, such as Google Cloud
  (GCP), if preferred.
• Browser: Chrome supports WebGPU,
  a powerful successor to WebGL, that
  brings the advancements of modern GPU APIs to the web.
• Automation: Puppeteer is a Node.js library that lets
  you control browsers programmatically with JavaScript. With Puppeteer, we can
  automate Chrome in headless mode, which means the browser runs without a
  visible interface, on a server. We are using the improved
  new headless mode, not the
  legacy form.

Verify the environment

The best way to check whether hardware acceleration is turned on in Chrome is to
type chrome://gpu into the address bar. You can
programmatically perform the equivalent with Puppeteer
with console.log or save the full report as PDF to check manually:

/* Incomplete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters: Expands later
const browser = await puppeteer.launch({
  headless: 'new',
  args:  ['--no-sandbox']
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({ path: './gpu.pdf' });

await browser.close();

Open chrome://gpu and you should have the following results:

Not a great start. It’s fairly clear that hardware detection was failing.
WebGL, WebGL2, and WebGPU are essentially disabled or software only. We
aren’t alone in this problem – there are numerous discussions online of people
in a similar situation, including on the official Chrome support channels
(1),
(2).

Enable WebGPU and WebGL support

By default, Headless Chrome
disables GPU.
To enable it on Linux, apply all of the following flags when launching Headless
Chrome:

• --no-sandbox flag disables Chrome’s security sandbox, which isolates the
  browser process from the rest of the system. Running Chrome as root without
  this sandbox is not supported.
• --headless=new flag runs Chrome with the new and improved
  headless mode, without any visible UI.
• --use-angle=vulkan flag tells Chrome to use the
  Vulkan backend
  for ANGLE, which
  translates OpenGL ES 2/3 calls to Vulkan API calls.
• --enable-features=Vulkan flag enables Vulkan graphics backend for
  compositing and rasterization in Chrome.
• --disable-vulkan-surface flag disables the VK_KHR_surface vulkan
  instance extension. Instead of using a swapchain,
  Bit blit is used for the
  present render result on screen.
• --enable-unsafe-webgpu flag enables the experimental WebGPU API in
  Chrome on Linux and disables the adapters blocklist.

Now we combine all the changes we have made so far. Here is the complete script.

/* Complete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters
const browser = await puppeteer.launch({
  headless: 'new',
  args: [
    '--no-sandbox',
    '--headless=new',
    '--use-angle=vulkan',
    '--enable-features=Vulkan',
    '--disable-vulkan-surface',
    '--enable-unsafe-webgpu',
  ]
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({path: './gpu.pdf'});

await browser.close();

Run the script again. No WebGPU problems are detected and the value changes from
disabled to software only.

However, hardware acceleration is still unavailable, the NVIDIA T4 GPU isn’t
detected.

Install the correct GPU drivers

We investigated more closely the output of chrome://gpu, with some GPU experts
on the Chrome team. We found issues with the default drivers installed on the
Linux Colab
instance, causing issues with Vulkan, leading to Chrome unable to detect the
NVIDIA T4 GPU at the GL_RENDERER level as shown in the following output. This
causes problems with Headless Chrome.

Installing the correct drivers that were compatible therefore fixes the issue.

To install the correct drivers, run the following commands during setup. The
last two lines help you to log the outputs of what NVIDIA drivers detects along
with vulkaninfo.

apt-get install -y vulkan-tools libnvidia-gl-525

// Verify the NVIDIA drivers detects along with vulkaninfo
nvidia-smi
vulkaninfo --summary

Now run the script again and we get the following result. 🎉

By using the correct drivers and flags when running Chrome, we now have WebGPU
and WebGL support using the shiny, new headless mode.

Behind the scenes: Our team’s investigation

After much research, we didn’t find working methods for the environment we
needed to execute in Google Colab, although there were some
hopeful posts
that worked in other environments, which was promising. Ultimately, we weren’t
able to replicate their success in the Colab NVIDIA T4 environment, as we had 2
key issues:

1. Some combinations of flags allow detection of the GPU, but don’t allow you to
   actually use the GPU.
2. Examples of working solutions by third parties used the old Chrome headless
   version, which at some point will be deprecated in favor of the
   new version. We needed a solution
   that worked with the new Headless Chrome to be better future proofed.

We confirmed the under utilization of the GPU by running an
example TensorFlow.js web page for image recognition,
whereby we trained a model to recognize clothing samples (sort of like a “hello
world” of machine learning).

On a regular machine, 50 training cycles (known as epochs) should run in less
than 1 second each. Calling Headless Chrome in its default state, we could log
the JavaScript console output to the Node.js server-side command line to see how
fast these training cycles were actually taking.

As expected, each training epoch took much longer than expected (several
seconds), which suggests Chrome has fallen back to plain old JS CPU execution
instead of utilizing the GPU:

After fixing the drivers and using the right combination of flags for Headless
Chrome, rerunning the TensorFlow.js training example results in much faster
training epochs.

Summary

Web AI has grown exponentially
since its creation in 2017. With browser technologies such as WebGPU, WebGL, and
WebAssembly, a machine learning model’s
mathematical operations can be further accelerated on the client side.

As of 2023 TensorFlow.js and MediaPipe Web crossed over 1 billion downloads of
models and libraries—a historic milestone and a sign of how web
developers and engineers are shifting to embrace AI in their next generation
web apps to make some truly incredible solutions.

With great success in usage comes great responsibility. At this level of usage
in production systems, the need arises for testing client-side, browser-based AI
models in a true browser environment, while also being scalable, automatable,
and within a known standardized hardware setup.

By harnessing the combined power of the new Headless Chrome and Puppeteer, you
can confidently test such workloads in a standardized and replicable
environment, ensuring consistent and reliable results.

Wrap up

A step-by-step guide is available in
our documentation, so you can try out the complete setup yourself.

If you found this useful, drop a shout out over on
LinkedIn,
X (formerly Twitter), or whatever
social network you use using hashtag #WebAI. It would be great to hear any
feedback you have so we know to write more stuff like this in the future.

Add a star on the Github repo
to receive any future updates.

Acknowledgements

A huge thank you to everyone on the Chrome team who helped debug the driver and
WebGPU issues we faced in this solution, with a special thanks to
Jecelyn Yeen and
Alexandra White for helping to wordsmith
this blog post. Thanks to Yuly Novikov, Andrey Kosyakov, and
Alex Rudenko who were instrumental
in creating the final, working solution.