ML Models#

import { Resources, SystemInfo } from "UnityEngine"
import { BackendType, ModelLoader, Tensor, TensorShape, Worker } from "Unity/InferenceEngine"
import type { ModelAsset } from "Unity/InferenceEngine"
import { Single } from "System"

const model = ModelLoader.Load(Resources.Load("mnist") as ModelAsset)
const worker = new Worker(model, BackendType.GPUCompute)

// Tensor<float>: a generic C# type is called with its type arguments.
const TensorFloat = Tensor(Single)
type FloatTensor = InstanceType<typeof TensorFloat>

const input = new TensorFloat(new TensorShape(1, 1, 28, 28), pixels, 0)
worker.Schedule(input)
const output = worker.PeekOutput() as FloatTensor

import { toArray } from "onejs-react"

const nextFrame = () => new Promise<void>((resolve) => requestAnimationFrame(() => resolve()))

output.ReadbackRequest()
while (!output.IsReadbackRequestDone()) await nextFrame()
const scores = toArray<number>(output.DownloadToArray())
input.Dispose()

const backend = SystemInfo.supportsComputeShaders ? BackendType.GPUCompute : BackendType.CPU

useEffect(() => {
    const worker = new Worker(ModelLoader.Load(asset), backend)
    return () => worker.Dispose()
}, [])

// mnist.ts
import { SystemInfo } from "UnityEngine"
import { BackendType, ModelLoader, Tensor, TensorShape, Worker } from "Unity/InferenceEngine"
import type { ModelAsset } from "Unity/InferenceEngine"
import { Single } from "System"
import { toArray } from "onejs-react"

// Sentis's Tensor<float>. OneJS binds a generic C# type by calling it with its
// type arguments.
const TensorFloat = Tensor(Single)
type FloatTensor = InstanceType<typeof TensorFloat>

const nextFrame = () => new Promise<void>((resolve) => requestAnimationFrame(() => resolve()))

/**
 * Compute shaders where the GPU has them, else the CPU. Never GPUPixel: on
 * WebGPU its readback is unsupported and the scores come back wrong.
 */
export function pickBackend(): BackendType {
    return SystemInfo.supportsComputeShaders ? BackendType.GPUCompute : BackendType.CPU
}

/** MNIST: 28 by 28 pixels in, one score per digit out. */
export class DigitClassifier {
    readonly backend: BackendType
    private worker: Worker
    private queue: Promise<unknown> = Promise.resolve()
    private disposed = false

    constructor(asset: ModelAsset, backend = pickBackend()) {
        this.backend = backend
        this.worker = new Worker(ModelLoader.Load(asset), backend)
    }

    /**
     * Ten scores for 784 pixels, row by row, 0 for background and 255 for ink.
     * Calls run one at a time, since each reads the worker's one output.
     */
    classify(pixels: Float32Array): Promise<number[]> {
        const run = this.queue.then(() => this.run(pixels))
        this.queue = run.catch(() => {})
        return run
    }

    dispose() {
        this.disposed = true
        this.worker.Dispose()
    }

    private async run(pixels: Float32Array): Promise<number[]> {
        if (this.disposed) throw new Error("classifier disposed")
        const input = new TensorFloat(new TensorShape(1, 1, 28, 28), pixels, 0)
        try {
            this.worker.Schedule(input)
            const output = this.worker.PeekOutput() as FloatTensor
            // A blocking download stalls the GPU, and WebGPU has none, so ask and wait.
            output.ReadbackRequest()
            while (!output.IsReadbackRequestDone()) {
                await nextFrame()
                if (this.disposed) throw new Error("classifier disposed")
            }
            return toArray<number>(output.DownloadToArray())
        } finally {
            input.Dispose()
        }
    }
}

/** Scores to probabilities. */
export function softmax(scores: number[]): number[] {
    const max = Math.max(...scores)
    const exps = scores.map((s) => Math.exp(s - max))
    const sum = exps.reduce((a, b) => a + b, 0)
    return exps.map((e) => e / sum)
}

// index.tsx (the model part)
import { Resources } from "UnityEngine"
import type { ModelAsset } from "Unity/InferenceEngine"
import { DigitClassifier, softmax } from "./mnist"

function App() {
    const [classifier, setClassifier] = useState<DigitClassifier | null>(null)
    const [probs, setProbs] = useState<number[] | null>(null)
    const pixels = useRef(new Float32Array(28 * 28))

    useEffect(() => {
        const c = new DigitClassifier(Resources.Load("mnist") as ModelAsset)
        let live = true
        c.classify(new Float32Array(28 * 28)).then(() => { if (live) setClassifier(c) })
        return () => {
            live = false
            c.dispose()
        }
    }, [])

    const predict = async () => {
        if (classifier) setProbs(softmax(await classifier.classify(pixels.current)))
    }
    // The pad: pointer events ink pixels.current with a soft brush, ink 255 on 0,
    // repaint it with useDrawing, and call predict() on pointer up.
}