使用 Transformers.js 製作 ML 驅動的網頁遊戲

釋出於 2023 年 7 月 5 日

在這篇博文中，我將向您展示我是如何製作塗鴉衝刺的，這是一個完全在瀏覽器中執行的即時 ML 驅動網頁遊戲（感謝 Transformers.js）。本教程的目標是向您展示製作自己的 ML 驅動網頁遊戲是多麼容易……正好趕上即將到來的開源 AI 遊戲大會（2023 年 7 月 7 日至 9 日）。如果您還沒有參加，請加入遊戲大會！

快速連結

演示： Doodle Dash
原始碼： doodle-dash
加入遊戲大會： Open Source AI Game Jam

概述

在開始之前，讓我們談談我們將要建立什麼。這款遊戲靈感來源於谷歌的 Quick, Draw! 遊戲，在該遊戲中，您會被給出一個單詞，神經網路有 20 秒的時間來猜測您正在畫什麼（重複 6 次）。事實上，我們將使用他們的訓練資料來訓練我們自己的草圖檢測模型！您難道不喜歡開源嗎？ 😍

在我們的版本中，您將有一分鐘的時間儘可能多地繪製物品，一次一個提示。如果模型預測出正確的標籤，畫布將被清除，您將獲得一個新單詞。一直這樣做，直到計時器用完！由於遊戲在您的瀏覽器中本地執行，我們根本不必擔心伺服器延遲。該模型能夠在您繪製時進行即時預測，每秒超過 60 次預測…… 🤯 太棒了！

本教程分為 3 個部分

訓練神經網路
使用 Transformers.js 在瀏覽器中執行
遊戲設計

1. 訓練神經網路

訓練資料

我們將使用谷歌 Quick, Draw! 資料集的子集來訓練我們的模型，該資料集包含 345 個類別的 500 多萬張圖畫。以下是資料集中的一些示例：

模型架構

我們將微調 `apple/mobilevit-small`，這是一個輕量級且適合移動裝置的視覺 Transformer，已在 ImageNet-1k 上進行預訓練。它只有 5.6M 個引數（檔案大小約 20MB），是瀏覽器內執行的理想選擇！有關更多資訊，請檢視 MobileViT 論文和下面的模型架構。

微調

為了使部落格文章（相對）簡短，我們準備了一個 Colab 筆記本，其中將向您展示我們對 `apple/mobilevit-small` 進行微調的精確步驟。總的來說，這包括

載入“Quick, Draw!”資料集。
使用`MobileViTImageProcessor`轉換資料集。
定義我們的排序函式和評估指標。
使用`MobileViTForImageClassification.from_pretrained`載入預訓練的 MobileVIT 模型。
使用`Trainer`和`TrainingArguments`輔助類訓練模型。
使用🤗 Evaluate評估模型。

注：您可以在 Hugging Face Hub 上此處找到我們微調過的模型。

2. 使用 Transformers.js 在瀏覽器中執行

什麼是 Transformers.js？

Transformers.js 是一個 JavaScript 庫，允許您直接在瀏覽器中執行 🤗 Transformers（無需伺服器）！它的設計旨在與 Python 庫功能等效，這意味著您可以使用非常相似的 API 執行相同的預訓練模型。

在幕後，Transformers.js 使用 ONNX Runtime，因此我們需要將我們微調過的 PyTorch 模型轉換為 ONNX。

將我們的模型轉換為 ONNX

幸運的是，🤗 Optimum 庫使得將您的微調模型轉換為 ONNX 變得超級簡單！最簡單（也是推薦的）方法是

克隆 Transformers.js 倉庫並安裝必要的依賴項

git clone https://github.com/xenova/transformers.js.git
cd transformers.js
pip install -r scripts/requirements.txt

執行轉換指令碼（它在底層使用 Optimum）
```
python -m scripts.convert --model_id <model_id>
```
其中 <model_id> 是您要轉換的模型名稱（例如 Xenova/quickdraw-mobilevit-small）。

設定我們的專案

讓我們首先使用 Vite 搭建一個簡單的 React 應用程式

npm create vite@latest doodle-dash -- --template react

接下來，進入專案目錄並安裝必要的依賴項

cd doodle-dash
npm install
npm install @xenova/transformers

然後可以透過執行以下命令啟動開發伺服器

npm run dev

在瀏覽器中執行模型

執行機器學習模型計算密集，因此在單獨的執行緒中執行推理非常重要。這樣我們就不會阻塞主執行緒，主執行緒用於渲染 UI 和響應您的繪圖手勢😉。 Web Workers API 使這變得超級簡單！

在 `src` 目錄中建立一個新檔案（例如，`worker.js`），並新增以下程式碼

import { pipeline, RawImage } from "@xenova/transformers";

const classifier = await pipeline("image-classification", 'Xenova/quickdraw-mobilevit-small', { quantized: false });

const image = await RawImage.read('https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/blog/ml-web-games/skateboard.png');

const output = await classifier(image.grayscale());
console.log(output);

我們現在可以透過將以下程式碼新增到 `App` 元件中，在 `App.jsx` 檔案中使用此 worker

import { useState, useEffect, useRef } from 'react'
// ... rest of the imports

function App() {
    // Create a reference to the worker object.
    const worker = useRef(null);

    // We use the `useEffect` hook to set up the worker as soon as the `App` component is mounted.
    useEffect(() => {
        if (!worker.current) {
            // Create the worker if it does not yet exist.
            worker.current = new Worker(new URL('./worker.js', import.meta.url), {
                type: 'module'
            });
        }

        // Create a callback function for messages from the worker thread.
        const onMessageReceived = (e) => { /* See code */ };

        // Attach the callback function as an event listener.
        worker.current.addEventListener('message', onMessageReceived);

        // Define a cleanup function for when the component is unmounted.
        return () => worker.current.removeEventListener('message', onMessageReceived);
    });

    // ... rest of the component
}

您可以透過執行開發伺服器（使用 `npm run dev`），訪問本地網站（通常是 https://:5173/），並開啟瀏覽器控制檯來測試一切是否正常工作。您應該會看到模型輸出被記錄到控制檯。

[{ label: "skateboard", score: 0.9980043172836304 }]

太棒了！🥳 儘管上述程式碼只是最終產品的一小部分，但它展示了機器學習方面的簡單性！其餘的只是使其看起來美觀並新增一些遊戲邏輯。

3. 遊戲設計

在本節中，我將簡要討論遊戲設計過程。提醒一下，您可以在 GitHub 上找到該專案的完整原始碼，因此我不會詳細介紹程式碼本身。

利用即時效能

在瀏覽器中進行推理的主要優點之一是我們可以即時進行預測（每秒超過 60 次）。在最初的 Quick, Draw! 遊戲中，模型每隔幾秒才進行一次新預測。我們可以在我們的遊戲中做同樣的事情，但那樣我們就無法利用其即時效能了！因此，我決定重新設計主遊戲迴圈

我們的版本不再是六個 20 秒的回合（每個回合對應一個新單詞），而是讓玩家在 60 秒內正確繪製儘可能多的塗鴉（一次一個提示）。
如果您遇到無法繪製的單詞，您可以跳過它（但這將花費您剩餘時間的 3 秒）。
在原版遊戲中，由於模型每隔幾秒鐘就會進行一次猜測，它會慢慢地從列表中劃掉標籤，直到最終猜對。在我們的版本中，我們反而會降低模型對前 n 個錯誤標籤的分數，其中 n 會隨著使用者持續繪圖而逐漸增加。