ONNX Runtime | Goodnotes 在 Windows、Web 和 Android 上利用 ONNX Runtime 实现涂写擦除功能

Scribble to Erase feature on Goodnotes for Windows, Web, and Android

在过去三年里，Goodnotes 工程团队一直在致力于将成功的 iPad 笔记应用带到其他平台，如 Windows、Web 和 Android。本文介绍这款 2022 年 iPad 年度应用是如何在三个平台上同时实现最受欢迎的 AI 功能之一——涂写擦除的，这得益于 ONNX Runtime。

📝 什么是涂写擦除？

我们都是人，所以都会犯错。涂写擦除是一个简单的功能，允许用户无需使用橡皮擦，只需在之前创建的内容上涂写即可删除内容。

Stylus scribbling to erase the word Abracadabra, then rabbit ears, then a hat

用户之前写的任何笔记，无论写了什么，都可以通过简单的涂写手势删除。这个功能对用户来说可能看起来很简单，但从工程角度来看却相当复杂。

Lorem Ipsum, with the ipsum scribbled out

这是 Goodnotes 工程团队在 Windows、Web 和 Android 上首次使用人工智能发布的功能，得益于 ONNX Runtime，它为边缘 AI 提供了跨平台的高性能模型推理。团队为此项目使用了内部训练的模型，并在三个不同平台上进行了设备端评估。

🔍 如何检测涂写？

对于 Goodnotes 来说，涂写手势不过是文档中添加的另一笔，遵循特殊的模式。要被视为涂写，一笔必须具备以下 2 个特征：

作为笔画一部分的点数必须足够多。
使用 ONNX Runtime 评估的 AI 模型应将这些笔画识别为涂写。

对于工程团队来说，这意味着对于 Goodnotes 团队添加到文档中的每一笔笔记，他们都必须评估其大小以及通过 AI 模型判断这新添加的一笔笔记是否是涂写。

一旦预处理阶段检查到笔记大小超过给定阈值，就该按照经典的 AI 模型评估流程进行，如下所示：

从点中提取笔记特征。
使用 ONNX Runtime 评估 AI 模型。

对于特征提取，Goodnotes 团队对笔记区域包含的点进行归一化，并将用户手写笔生成的点列表转换为浮点数组。这个过程无非是所有 AI 模型遵循的经典特征提取流程，目的是将用户数据转换为 AI 模型能够理解的形式。

此项目使用的 AI 模型是一个基于 LSTM 的监督模型，团队对其进行了精心设计，并将其部署到所有平台，以便所有用户即使在没有网络连接的情况下也能在设备上进行评估。

一旦点被表示为 AI 模型可以处理的形式，通过使用 ONNX Runtime 并读取模型输出作为分数，我们就可以判断最近添加的笔记是否是涂写。如果该笔画被认为是涂写，则其下方所有笔记都将自动删除。

Diagram of ONNX Runtime workflow for handwritten note recognition

🤝 为什么选择 ONNX Runtime？

当 Goodnotes 团队必须评估此功能的实现时，他们需要做出一个决定：如何评估 AI 模型。这是该项目首次使用 AI，并且该产品的 iOS 版本使用的是 CoreML，这与当前的项目技术栈不兼容，因为此 Apple 技术在 iOS/MacOS SDK 之外不可用。所以他们决定尝试一些不同的东西。

Goodnotes 在 Windows、Web 和 Android 上的技术栈基于 Web 技术。底层应用使用了 Progressive Web Application（渐进式 Web 应用）。当应用从 Microsoft Store 或其他商店（如 Google Play）安装时，应用使用原生包装器，但最终，这个项目是一个作为全屏原生应用运行的 Web 应用。这意味着用于评估 AI 模型的技术必须与 Web 技术栈兼容，并且还必须足够高性能以满足他们的需求，并在可能的情况下启用硬件运行时。因此，在检查不同替代方案时，团队发现 ONNX 作为一种可移植格式，以及带有 Web 解决方案的 ONNX Runtime，决定试一试。在进行了几次实验并在此功能实现之前使用 ONNX Runtime 创建了一些原型后，团队决定这是正确的技术选择！

团队选择使用 ONNX Runtime 而非其他技术有以下四个原因：

开发的原型证明了 ONNX Runtime 集成对我们来说相当简单，并提供了我们所需的所有功能。
ONNX 是一种可移植格式，我们可以用它将当前的 CoreML 模型导出为一种可以在许多不同操作系统上评估的格式。
ONNX Runtime 中的执行提供程序（Execution providers）提供了硬件特定的加速，这使得我们在评估模型时能够获得最佳性能。特别是 Web 解决方案，它有针对 CPU 执行的 WSAM 执行提供程序，以及用于进一步加速（利用 GPU/NPU）的 WebNN 和 WebGPU 执行提供程序，这些对我们来说都是非常有趣的例子。
与 AI 模型的 LSTM 设计兼容。

💻 我们的 ONNX Runtime 代码看起来如何？

Goodnotes 团队与 iOS/Mac 团队共享 Goodnotes 应用的大部分业务逻辑代码。这意味着他们编译原始的 Swift 代码库，处理笔画，并通过 Web Assembly 处理 Swift 输出的模型结果。但在执行栈中有一个点，Goodnotes 团队必须评估模型，他们在这里使用 ONNX Runtime 将执行从 Swift 委托给 Web 环境。

评估模型的 TypeScript ONNX Runtime 代码类似于以下代码片段：

export class OnnxScribbleToEraseAIModel extends OnnxAIModel<Array<Array<number>>, EvaluationResult>

implements ScribbleToEraseAIModel
{
    getModelResource(): OnDemandResource {
        return OnDemandResource.ScribbleToErase;
    }


    async evaluateModel(input: Array<Array<number>>): Promise<EvaluationResult> {
        const startTime = performance.now();
        const { tensor, initializeTensorTime } = this.initializeTensor(input);
        const { evaluationScore, evaluateModelTime } = await this.runModel(tensor);
        const result = {
            score: evaluationScore ?? 0.0,
            timeToInitializeTensor: initializeTensorTime,
            timeToEvaluateTheModel: evaluateModelTime,
            totalExecutionTime: performance.now() - startTime,
        };
        return result;
    }
…..

正如你所见，实现是你从任何 AI 功能中期望的经典代码。输入数据作为特征数组获取，我们稍后使用张量将其馈送给模型。评估完成后，我们检查作为模型输出获得的分数，如果分数高于特定阈值，则认为输入是涂写。

正如你在代码中看到的，除了初始化张量和评估模型之外，我们还在跟踪执行时间，以便验证我们的实现，并更好地了解实际用户使用此功能时在生产中所需的资源。

private initializeTensor(input: number[][]) {
       const prepareTensorStartTime = performance.now();
       const modelInput = new Float32Array(input.flat());
       const tensor = new Tensor(modelInputTensorType, modelInput, modelInputDimensions);
       const initializeTensorTime = performance.now() - prepareTensorStartTime;
       return { tensor, initializeTensorTime };
   }


private async runModel(tensor: Tensor) {
    const evaluateModelStartTime = performance.now();
    const inferenceSession = this.session;
    const outputMap = await inferenceSession.run({ x: tensor });
    const outputTensor = outputMap[modelOutputName];
    const evaluationScore = outputTensor?.data[0] as number | undefined;
    const evaluateModelTime = performance.now() - evaluateModelStartTime;
    return { evaluationScore, evaluateModelTime };
}

最重要的是，在此案例中，Goodnotes 团队决定使用 ONNX Runtime 从 Web Worker 加载和评估 AI 模型，并在 Web Worker 中运行推理会话，因为此路径在我们的应用中处于关键的用户体验流程中，我们希望最大程度地减少对用户性能的影响。

ort.env.logLevel = 'fatal';
ort.env.wasm.wasmPaths = '/onnx/providers/wasm/';
this.session = await InferenceSession.create(modelURL);

根据模型架构，此项目配置的执行提供程序是 CPU 提供程序。这是一个轻量级模型，使用底层由 WASM 支持的默认 CPU 执行提供程序，我们可以获得相当快的执行时间。我们计划在新的 AI 场景中，对更高级的模型使用 WebGPU 和 WebNN 执行提供程序。

🚀 部署与集成

由于团队使用的技术栈，Web 技术的使用使得 ONNX Runtime 的集成以及我们托管 AI 模型的方式值得一提。对于这个项目，Goodnotes 使用 Vite 作为前端工具，因此他们不得不稍微修改 Vite 配置，不仅要分发我们的 AI 模型，还要分发 CPU 执行提供程序所需的资源。这对于团队来说不是什么大问题，因为 ONNX Runtime 文档已经涵盖了打包器的使用，但这很有趣，因为该应用是一个可以离线使用的 PWA，此更改增加了捆绑包的大小，不仅包含了模型二进制文件，还包含了 ONNX Runtime 所需的所有资源。

📈 投产数月后的结果

Goodnotes 数月前发布了此功能。从第一天起，所有用户就开始透明地使用这个涂写擦除模型。有些人主动开始使用涂写手势删除内容，而另一些人则将此功能发现为一种自然的手势。

自发布日期以来，Goodnotes 团队使用 ONNX Runtime 评估 AI 模型已将近 20 亿次！使用 CPU 执行提供程序，并在 Worker 中运行模型，团队获得了 P95 的评估时间低于 16 毫秒，P99 低于 27 毫秒！来自世界各地、不同操作系统和平台的用户已经使用涂写擦除功能修改了他们的笔记，团队对通过这款出色的 ONNX Runtime 解决方案所取得的技术成就感到非常自豪！