Native Model Authoring for Edge with LiteRT Tensor API

通过 LiteRT 张量 API 在边缘设备上原生构建模型 - Google 开发者博客

Google Tag Manager (noscript)

结束 Google Tag Manager (noscript)

HTML

通过 LiteRT 张量 API 在边缘设备上原生构建模型

2026 年 10 月 19 日

Ping Yu

软件工程师

Quentin Khan

高级软件工程师

Chintan Parikh

高级产品经理

Lu Wang

高级软件工程师

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

在设备上部署 AI 需要满足针对每个目标设备的严格延迟和内存预算。目前生产环境中部署模型最常见的路径是将训练好的模型进行转换，例如使用 LiteRT Torch 将 PyTorch 模型转换，然后通过 LiteRT 运行时执行。这种方式能提供最快且最可靠的生产部署路径，也是我们推荐给大多数项目的方案。不过，有时开发者需要对网络的表示和优化方式拥有完全控制权，以便支持新架构或尝试新的降级逻辑，而无需等待上游框架更新。

为了获得终极控制权，开发者需要手动使用原生 C++ 或着色器语言编写计算。Apple 的 MLX 在 Apple 芯片上使这种以代码为核心的开发方式流行起来，而 llama.cpp 则展示了定制化 C++ 实现对大语言模型推理的提升潜力。然而，这种控制方式需要付出显著代价：实现与特定后端（如 Metal、OpenCL 或 Vulkan）紧密耦合，每次面对新硬件目标都需要重写代码。此外，在 CPU、GPU 和 NPU 之间管理零拷贝缓冲区互操作性和同步，同时将预处理、推理和后处理环节拼接在一起，需要大量手动工程工作。

为了解决上述挑战，我们非常高兴地宣布推出全新的 LiteRT 张量 API，这是 LiteRT 堆栈中一个轻量级的模型构建层，旨在提升手动构建模型的效率。该 API 将原生开发的灵活性与跨平台可移植性相结合。你可以在 C++ 中逐个操作构建计算图，精确选择每个操作的执行方式及其连接方式。相同的图结构可以在 LiteRT 的 CPU、GPU（由开源 ML Drift 支持）和 NPU 上运行，或通过 XNNPACK 直接在 CPU 上运行而无需模型文件或运行时，因此无需为每个新硬件目标重写代码。模型推理与预处理/后处理可以通过相同 API 组装，并通过共享加速器缓冲区串联，使数据无需往返 CPU。建议从转换路径开始，这仍然是大多数模型的最佳默认方案。当需要对图结构进行底层控制，或希望整个流水线在加速器上端到端运行时，再使用张量 API。

LiteRT 张量 API 同时也为 AI 编程代理提供了理想的程序化基础，通过结构化原语可自动生成、重新连接和调整模型图。在本文的其余部分，我们将深入探讨其设计，并通过实际代码示例演示如何从零开始构建模型。

实际应用案例：在 Web 上运行 Segment Anything

抱歉，你的浏览器不支持播放此视频

为了展示张量 API 的实际应用，我们构建了一个交互式 Segment Anything 应用，完全在浏览器的 WebGPU 上运行。该应用通过 Meta 的 Segment Anything Model 2（SAM 2）实现实时多对象分割和跟踪，用户可通过简单点击或由 Gemma 4 E2B 提供支持的自然语言提示触发。

LiteRT Tensor API 是用 C++ 编写的 SAM 2 模型的完整实现。一个编码代理通过 WebGPU 将其从每秒不到 10 帧优化到 512x512 分辨率下超过 30 帧。LiteRT Samples GitHub 仓库中已开源完整实现和网页应用。

LiteRT Tensor API 是什么？

LiteRT Tensor API 为设备端模型的编写和优化提供四个关键特性：

• 对图结构的完全控制：你可以决定模型图中的每个操作以及它们之间的所有连接，从逐元素数学运算到加速器优化的 LLM 复合操作（如缩放点积注意力、RMSNorm 和 KV 缓存更新）。

• 直接权重导入：可直接从模型检查点（如 Hugging Face .safetensors 文件）加载权重，无需中间转换脚本。

• 以张量为中心的缓冲区互操作：预处理、模型分区和后处理通过共享加速器缓冲区传递张量，使中间数据在加速器上保持零拷贝。

• AI 代理的构建模块：模型由明确定义的操作构建的纯 C++ 代码组成，因此编码代理可以构建模型、将其结果与参考模型对比，并通过在每台设备上对变体进行基准测试来调整性能。

跨平台性能

我们使用相同的 C++ 代码在三台设备的 GPU 上运行 SAM 2 模型：一台 macOS 笔记本（Metal）、一部 iOS 手机（Metal）和一部 Android 手机（OpenCL）。在 MacBook Pro 的测试中，其运行速度比 Apple MLX（fp16，图像编码器加掩码解码器）快约 1.5 倍。SAM 2.1 示例包含基准测试命令，你可以复现这些结果。

注意：时间涵盖图像编码器和掩码解码器；预处理和后处理增加约 0.34 毫秒。输出掩码与 Meta 的 PyTorch 参考实现逐像素匹配。

从原生 C++ 开始使用 Tensor API

用原生 C++ 编写模型

通过 Tensor API，每一层都是基于张量值的 C++ 函数，由高级操作符和直接从检查点加载的权重构建。例如，这是一个多头自注意力层的实现：

// 对 x: [batch, tokens, dim] 进行多头自注意力。
Tensor Attention(const Tensor& x, const WeightMap& w, int heads) {
  std::vector<int> s = x.GetShape();
  int b = s[0], n = s[1], dim = s[2], d = dim / heads;
  Tensor q = FullyConnected(x, w.at("q_proj.weight"), w.at("q_proj.bias"));
  Tensor k = FullyConnected(x, w.at("k_proj.weight"), w.at("k_proj.bias"));
  Tensor v = FullyConnected(x, w.at("v_proj.weight"), w.at("v_proj.bias"));
  // 分割头：[b, n, dim] -> [b, heads, n, d].
  q = Transpose(Reshape(q, {b, n, heads, d}), {0, 2, 1, 3});
  k = Transpose(Reshape(k, {b, n, heads, d}), {0, 2, 1, 3});
  v = Transpose(Reshape(v, {b, n, heads, d}), {0, 2, 1, 3});
  // 对所有头同时执行 softmax(q * k^T / sqrt(d)) * v.
  Tensor scores = BatchMatMul(q, k, /*adj_x=*/false, /*adj_y=*/true);
  Tensor attn = Softmax(Mul(scores, 1.0f / std::sqrt(static_cast<float>(d))));
  Tensor out = BatchMatMul(attn, v);
  // 合并头：[b, heads, n, d] -> [b, n, dim]，然后投影。
  out = Reshape(Transpose(out, {0, 2, 1, 3}), {b, n, dim});
  return FullyConnected(out, w.at("out_proj.weight"), w.at("out_proj.bias"));
}

C

复制

在 SAM 2 示例中，BuildEncoder 和 BuildDecoder 通过这种方式将层组合成编码器和解码器图，ModelFactory 将两者导出为 sam2_image.tflite 文件的两个签名。

.tflite 文件是运行图的两种方式之一。后端是 Tensor 的模板参数：TfLiteMixinTag 导出的 .tflite 文件由 LiteRT 在 CPU、ML Drift（GPU）或 NPU 上运行，而 XnnpackMixinTag 则通过 XNNPACK 直接在 CPU 上运行相同代码，无需模型文件。同时标记两个标签以保留两种选项：

// 为两个后端标记张量；操作记录任一后端均可运行的图。
using T = Tensor<TfLiteMixinTag, XnnpackMixinTag>;
T y = Attention(x, w, /*heads=*/8);  // 上层操作，作用于 T 类型值。
// 导出 .tflite 文件，并通过 LiteRT 在 CPU、GPU 或 NPU 上运行。
Save({y}, "attention.tflite");
// 或者通过 XNNPACK 直接在 CPU 上运行相同图。
auto runner = XnnpackRunner::Create({y});
runner->Run();

SAM 2 通过 .tflite 路径在浏览器中实现 WebGPU 支持，本文其余部分也将沿用此路径。Gemma 4 示例则通过直接路径在 CPU 上执行预填充和解码，这是运行 Tensor API 代码最轻量的方式。Tensor API README 文档会逐步演示这两种方式。

自然的工作流程会同时使用两者：在直接 XNNPACK 路径上迭代，直接从 .safetensors 检查点加载权重，无需重建模型文件；当准备在移动端或通过 LiteRT.js 在网页端部署时，再导出 .tflite 文件。

表达完整的机器学习计算图

实际应用需要的不只是神经网络本身。在演示中，将摄像头画面和点击转换为掩码需要四个步骤：预处理将画面转换为模型预期的 512x512 图像；图像编码器将图像转换为特征；掩码解码器将这些特征与点击结合以预测掩码；后处理将掩码缩放回画面尺寸。要快速运行整个流水线，需要两个关键点：在步骤之间保持数据在加速器（如 GPU）上，以及控制每个步骤在每个加速器上的运行方式。

零拷贝：将整个流水线保留在加速器上

在典型应用中，只有模型在 GPU 上运行；CPU 负责调整每个摄像头画面尺寸，将编码器的输出传递给解码器，并对掩码进行后处理。每次数据传输都会复制数据并阻塞 GPU。

在 litert-samples/sam2 中，所有四个步骤都在 GPU 上作为 Tensor API 图运行。例如，预处理是一个使用 CreateLambdaRunner 构建的小型图。这些图在启动时一次性连接以共享缓冲区，然后为每个画面和点击运行：

// 预处理使用与模型相同的 Tensor API 操作编写：
// RGBA 摄像头画面 [1, H, W, 4] -> 标准化像素 [1, 512, 512, 3].
auto pre_runner = CreateLambdaRunner(
    env, options,
    {{"frame", Tensor({.type = Type::kU8, .shape = {1, H, W, 4}})}},
    [=](const TensorsMap& inputs) {
      Tensor rgb = Slice(inputs.at("frame"), {0, 0, 0, 0}, {1, H, W, 3});
      Tensor resized = ResizeBilinear(Cast(rgb, Type::kFP32), {512, 512});
      return TensorsMap{{"pixels", Add(Mul(resized, scale), bias)}};
    });
auto model = LitertDynamicRunner::Create(env, "sam2_image.tflite", options);
// 一次性连接：相邻步骤共享缓冲区，无需复制数据。
pre_runner.SetOutput("pixels", model.GetInput("encode_image", "pixels"));
for (auto name : {"image_embeddings", "feat_s1", "feat_s0"}) {
  model.SetInput("decode_mask", name, model.GetOutput("encode_image", name));
}
post_runner.SetInput("masks", model.GetOutput("decode_mask", "masks"));

无数据返回到 CPU：所有 522 个编码器和 269 个解码器操作均在 GPU 上运行，预处理和后处理每帧仅增加约 0.34 毫秒。

完全控制：为每个后端调整图结构

由于图结构是纯 C++ 编写的，你可以或你的编码代理可以为每个后端重新编写它。我们的代理 Antigravity 将 SAM 2 的每帧决策（例如对象是否仍然可见）从 CPU 代码转移到了图结构中。像 score > 0 这样的比较在 WebGPU 上仍会回退到 CPU，因为它们输出 true/false 值，因此代理在 StepPos 中将它们改写为算术运算：

// 简化自 sam2v_graph.cc 中的 StepPos 和 BuildStep。
// 之前：Greater 输出 true/false 值，会回退到 CPU。
// Tensor visible = Cast(Greater(score, zero), Type::kFP32);
// 之后：GPU 友好操作；score <= 0 时为 0，score >= 1e-4 时为 1。
Tensor visible = Sub(1.0f, Relu(Sub(1.0f, Mul(Relu(score), 1e4f))));
// 如果对象可见则保留掩码，否则将其清空。
Tensor mask = Add(Mul(best_mask, visible), Mul(Sub(1.0f, visible), -1024.0f));

现在每个对象每帧只需一次 GPU 调用。代理进行了类似的一系列修改，例如将注意力缩放合并到权重中，并在保留之前与 Hugging Face 的 SAM 2 进行对比验证。正是通过这些优化，我们使模型的帧率提升了三倍以上。

后续计划

我们正在开发直接 ML Drift（GPU）后端、NPU 友好操作以及自动全图量化。我们还将添加 Python 编写支持（支持导出为 .tflite），并扩展对动态形状的支持，以适应不同图像尺寸、批量大小和序列长度。

探索我们的资源，开始你的 LiteRT 之旅：

• 使用编码代理？指向 Tensor API README 和 SAM 2 示例；它们共同涵盖了 API 和一个可构建的端到端示例。

• 文档：LiteRT Tensor API 指南（入门指南、构建说明和示例）和 Tensor API 参考文档。

• 参考实现：使用 LiteRT Tensor API 用 C++ 编写的 SAM 2，Live Web Demo 以及 Web 应用源码。

• LLM 示例：Gemma 4 直接在 XNNPACK（CPU）上运行，Gemma 3 通过 LiteRT（CPU 和 GPU）运行，两者均使用 Tensor API 编写。

• 模型定义：sam2_graph.cc - Hiera 编码器和掩码解码器图结构。

致谢

LiteRT Tensor API 及其参考实现是 Google AI Edge 和 LiteRT 团队协作努力的成果。特别感谢 Byungchul Kim、Changming Sun、Chenchen Tang、Dillon Sharlet、Jingjiang Li、Lu Wang、Lutz Justen、Matthias Grundmann、Misha Gutman、Ronald Wotzlaw 和 Sachin Kotwani。

发布于：

• 移动端

• 网络

• 人工智能

• 案例研究

• 公告

• 活动

• 学习

• 探索

• 解决方案

• 获得灵感

• 资源

上一篇

下一篇

相关文章

列表

人工智能

案例研究

公告

2026 年 10 月 13 日

移动端

网络

操作指南

宣布 Google OAuth 支持 DPoP：保护刷新令牌

2026 年 10 月 15 日

ML Drift：下一代边缘 GPU AI/ML 推理

2026 年 10 月 8 日

导航点