Turbine os testes do modelo de IA da Web: WebGPU, WebGL e Headless Chrome

Jason Mayes

François Beaufort

Temos ótimas notícias! Você criou um aplicativo de IA da Web legal que executa modelos de machine learning diretamente no dispositivo de um usuário. Ele é executado inteiramente no navegador da Web do cliente, sem depender da nuvem. Esse design no dispositivo melhora a privacidade do usuário, aumenta o desempenho e reduz significativamente os custos.

No entanto, há um obstáculo. Seu modelo do TensorFlow.js pode operar em CPUs (WebAssembly) e GPUs mais potentes (usando o WebGL e o WebGPU). A pergunta é: como é possível automatizar de forma consistente os testes de navegador com o hardware selecionado?

Manter a consistência é fundamental para comparar o desempenho do modelo de aprendizado de máquina ao longo do tempo, à medida que você itera e melhora, antes da implantação para que os usuários reais usem no dispositivo.

Configurar um ambiente de teste consistente com GPUs pode ser mais difícil do que esperado. Nesta postagem do blog, vamos compartilhar os problemas que enfrentamos e como os resolvemos para que você possa melhorar o desempenho do seu aplicativo.

Não é apenas para desenvolvedores de IA da Web. Se você trabalha com jogos ou gráficos da Web, esta postagem também é valiosa para você.

O que há na nossa caixa de ferramentas de automação

Confira o que estamos usando:

• Ambiente: um notebook do Google Colab baseado em Linux conectado a uma GPU NVIDIA T4 ou V100. Você pode usar outras plataformas de nuvem, como o Google Cloud (GCP), se preferir.
• Navegador: o Chrome oferece suporte à WebGPU, um poderoso sucessor da WebGL que traz os avanços das APIs de GPU modernas para a Web.
• Automação: o Puppeteer é uma biblioteca do Node.js que permite controlar navegadores de forma programática com JavaScript. Com o Puppeteer, podemos automatizar o Chrome no modo headless, o que significa que o navegador é executado sem uma interface visível em um servidor. Estamos usando o novo modo headless aprimorado, não o formulário legado.

Verificar o ambiente

A melhor maneira de verificar se a aceleração de hardware está ativada no Chrome é digitar chrome://gpu na barra de endereço. É possível executar o equivalente com o Puppeteer com console.log ou salvar o relatório completo como PDF para verificar manualmente:

/* 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();

Abra chrome://gpu e confira os resultados:

Não é um bom começo. É bastante claro que a detecção de hardware estava falhando. A WebGL, a WebGL2 e a WebGPU são essencialmente desativadas ou somente software. Não somos os únicos com esse problema. Há muitas discussões on-line de pessoas em uma situação semelhante, inclusive nos canais oficiais de suporte do Chrome (1), (2).

Ativar o suporte a WebGPU e WebGL

Por padrão, o Chrome Headless desativa a GPU. Para ativar no Linux, aplique todas as flags a seguir ao iniciar o Chrome sem cabeça:

• A flag --no-sandbox desativa o sandbox de segurança do Chrome, que isola o processo do navegador do restante do sistema. Não é possível executar o Chrome como raiz sem este sandbox.
• A flag --headless=new executa o Chrome com o novo e melhorado modo headless, sem interface visível.
• A flag --use-angle=vulkan instrui o Chrome a usar o back-end do Vulkan para o ANGLE, que converte chamadas OpenGL ES 2/3 em chamadas da API Vulkan.
• A flag --enable-features=Vulkan ativa o back-end gráfico do Vulkan para composição e rasterização no Chrome.
• A flag --disable-vulkan-surface desativa a extensão de instância do vulkan VK_KHR_surface. Em vez de usar uma cadeia de troca, o Blit de bit é usado para apresentar o resultado da renderização na tela.
• A flag --enable-unsafe-webgpu ativa a API WebGPU experimental no Chrome no Linux e desativa a lista de bloqueio de adaptadores.

Agora, vamos combinar todas as mudanças feitas até agora. Confira o script completo.

/* 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();

Execute o script novamente. Nenhum problema com a WebGPU é detectado, e o valor muda de desativado para somente software.

No entanto, a aceleração de hardware ainda não está disponível, a GPU NVIDIA T4 não é detectada.

Instalar os drivers corretos da GPU

Investigamos mais de perto a saída de chrome://gpu com alguns especialistas em GPU da equipe do Chrome. Encontramos problemas com os drivers padrão instalados na instância do Linux Colab, o que causou problemas com o Vulkan, impedindo que o Chrome detectasse a GPU NVIDIA T4 no nível GL_RENDERER, conforme mostrado na saída a seguir. Isso causa problemas com o Headless Chrome.

A instalação dos drivers corretos que eram compatíveis corrigiu o problema.

Para instalar os drivers corretos, execute os comandos abaixo durante a configuração. As duas últimas linhas ajudam a registrar as saídas do que os drivers da NVIDIA detectam com vulkaninfo.

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

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

Agora execute o script novamente e confira o resultado a seguir. 🎉

Ao usar os drivers e flags corretos ao executar o Chrome, agora temos suporte para WebGPU e WebGL usando o novo modo headless.

Bastidores: a investigação da nossa equipe

Depois de muita pesquisa, não encontramos métodos que funcionassem para o ambiente que precisávamos executar no Google Colab, embora houvesse algumas postagens promissoras que funcionaram em outros ambientes. Não conseguimos replicar o sucesso no ambiente Colab NVIDIA T4, porque tivemos dois problemas principais:

1. Algumas combinações de flags permitem a detecção da GPU, mas não permitem que você use a GPU.
2. Exemplos de soluções de trabalho de terceiros usaram a versão antiga do Chrome headless, que em algum momento será descontinuada em favor da nova versão. Precisávamos de uma solução que funcionasse com o novo Chrome Headless para estar mais preparado para o futuro.

Confirmamos a subutilização da GPU executando um exemplo de página da Web do TensorFlow.js para reconhecimento de imagem, em que treinamos um modelo para reconhecer amostras de roupas (mais ou menos como um "hello world" de machine learning).

Em uma máquina normal, 50 ciclos de treinamento (conhecidos como épocas) precisam ser executados em menos de 1 segundo cada. Ao chamar o Chrome sem cabeça no estado padrão, foi possível registrar a saída do console JavaScript na linha de comando do lado do servidor do Node.js para saber a velocidade desses ciclos de treinamento.

Como esperado, cada período de treinamento levou muito mais tempo do que o esperado (vários segundos), o que sugere que o Chrome voltou à execução da CPU do JS, em vez de utilizar a GPU:

Depois de corrigir os drivers e usar a combinação certa de flags para o Chrome Headless, a reexecução do exemplo de treinamento do TensorFlow.js resulta em épocas de treinamento muito mais rápidas.

Resumo

A IA da Web cresceu exponencialmente desde a criação, em 2017. Com tecnologias de navegador, como WebGPU, WebGL e WebAssembly, as operações matemáticas de um modelo de machine learning podem ser aceleradas ainda mais no lado do cliente.

Em 2023, o TensorFlow.js e o MediaPipe Web ultrapassaram 1 bilhão de downloads de modelos e bibliotecas, um marco histórico e um sinal de como desenvolvedores e engenheiros da Web estão mudando para adotar a IA nos apps da Web de última geração para criar soluções realmente incríveis.

Com o sucesso no uso, vem uma grande responsabilidade. Nesse nível de uso em sistemas de produção, surge a necessidade de testar modelos de IA baseados em navegador do lado do cliente em um ambiente de navegador real, além de serem escalonáveis, automatizáveis e dentro de uma configuração de hardware padronizada conhecida.

Ao aproveitar o poder combinado do novo Headless Chrome e do Puppeteer, você pode testar essas cargas de trabalho em um ambiente padronizado e replicável, garantindo resultados consistentes e confiáveis.

Conclusão

Um guia detalhado está disponível na nossa documentação para que você possa testar a configuração completa.

Se você achou este artigo útil, deixe um comentário no LinkedIn, X (antigo Twitter) ou em qualquer rede social que você usa usando a hashtag #WebAI. Seria ótimo receber seu feedback para que possamos escrever mais artigos como este no futuro.

Adicione uma estrela no repositório do GitHub para receber atualizações futuras.

Agradecimentos

Agradecemos muito a todos da equipe do Chrome que ajudaram a depurar o driver e os problemas da WebGPU que enfrentamos nessa solução, com um agradecimento especial a Jecelyn Yeen e Alexandra White por ajudar a escrever este post do blog. Agradecemos a Yuly Novikov, Andrey Kosyakov e Alex Rudenko, que foram fundamentais na criação da solução final.