Supercharge Web AI model testing: WebGPU, WebGL, and Headless Chrome

Jason Mayes

François Beaufort

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.