Google AI Edge Torch的特性詳解

作者 / 主任工程師 Cormac Brick，軟件工程師 Advait Jain，軟件工程師 Haoliang Zhang

我們很高興地發(fā)布 Google AI Edge Torch，可將 PyTorch 編寫的模型直接轉(zhuǎn)換成 TFLite 格式 (.tflite)，且有著優(yōu)異的模型覆蓋率和 CPU 性能。TFLite 已經(jīng)支持 Jax、Keras 和 TensorFlow 編寫的模型，現(xiàn)在我們加入了對(duì) PyTorch 的支持，進(jìn)一步豐富了框架選擇。

這一新產(chǎn)品現(xiàn)已作為 Google AI Edge 的一部分提供。Google AI Edge 是一套易于使用的工具，包含可直接使用的機(jī)器學(xué)習(xí) (ML) 任務(wù)、構(gòu)建機(jī)器學(xué)習(xí)流水線的框架，以及運(yùn)行流行的大語言模型 (LLM) 和自定義模型的能力——所有這些都可在設(shè)備上運(yùn)行。本文是 Google AI Edge 博客連載中的第一篇，用于幫助開發(fā)者們構(gòu)建 AI 功能，并輕松地將其部署至多個(gè)平臺(tái)。

今天發(fā)布的 AI Edge Torch Beta 版本包含以下特性:

直接集成 PyTorch

出色的 CPU 性能和初步 GPU 支持

在 torchvision、timm、torchaudio 和 HuggingFace 里的 70 多個(gè)模型上得到驗(yàn)證

支持超過 70% 的 PyTorch core_aten 算子

兼容現(xiàn)有的 TFLite 運(yùn)行時(shí)，無需更改部署代碼

支持在工作流的多個(gè)階段進(jìn)行模型探索器 (Model Explorer) 可視化

以 PyTorch 為中心的簡(jiǎn)潔體驗(yàn)

Google AI Edge Torch 從一開始就致力于為 PyTorch 社區(qū)提供卓越的開發(fā)體驗(yàn)，API 使用起來感覺非常原生，并提供簡(jiǎn)便的模型轉(zhuǎn)換路徑。

import torchvision
import ai_edge_torch


# Initialize model
resnet18 = torchvision.models.resnet18().eval()


# Convert
sample_input = (torch.randn(4, 3, 224, 224),)
edge_model = ai_edge_torch.convert(resnet18, sample_input)


# Inference in Python
output = edge_model(*sample_input)


# Export to a TfLite model for on-device deployment
edge_model.export('resnet.tflite'))

在底層，ai_edge_torch.convert()使用 torch.export 集成了 TorchDynamo——在 PyTorch 2.x 中，這個(gè)方法用于將 PyTorch 模型導(dǎo)出為標(biāo)準(zhǔn)化的模型形式，從而在不同環(huán)境中運(yùn)行。我們目前的實(shí)現(xiàn)支持超過 70% 的 core_aten 算子，這個(gè)比例會(huì)在構(gòu)建 ai_edge_torch 1.0 版本的過程中大幅增加。我們還提供了 PT2E 量化的示例，這是 PyTorch2 原生的量化方法，以簡(jiǎn)化量化工作的流程。我們很期待聽到來自 PyTorch 社區(qū)的反饋，以進(jìn)一步改善開發(fā)者體驗(yàn)，從而幫助大家更好地把用 PyTorch 打造的新穎體驗(yàn)部署至更多樣的設(shè)備中。

模型覆蓋和性能

在此版本發(fā)布之前，許多開發(fā)者使用社區(qū)提供的轉(zhuǎn)換方法，如 ONNX2TF，在 TFLite 中運(yùn)行 PyTorch 模型。我們開發(fā) AI Edge Torch 的目標(biāo)是減少開發(fā)過程中的阻力，提供出色的模型覆蓋率，并繼續(xù)完成我們的使命: 在 Android 設(shè)備上提供最佳的性能。

在覆蓋率方面，我們的測(cè)試表明，與現(xiàn)有工作流程 (尤其是 ONNX2TF) 相比，AI Edge Torch 在給定的模型集合上的覆蓋率有顯著的提高。

在性能方面，我們的測(cè)試顯示 AI Edge Torch 與 ONNX2TF 的基準(zhǔn)性能表現(xiàn)相當(dāng)，比 ONNX 運(yùn)行時(shí)相比則有著更好的性能。

下圖顯示了在 ONNX 覆蓋的模型子集上的每個(gè)模型的詳細(xì)性能:

△ 相對(duì)于 ONNX 的每個(gè)網(wǎng)絡(luò)的推理延遲。以 Pixel 8 為測(cè)試設(shè)備，使用 fp32 精度。XNNPACK 固定為 4 個(gè)線程以確保復(fù)現(xiàn)性，經(jīng)過 20 次迭代預(yù)熱后取 100 次運(yùn)行的均值

早期體驗(yàn)用戶和合作伙伴

在過去的幾個(gè)月中，我們與參與早期體驗(yàn)的合作伙伴們密切合作，包括 Shopify、Adobe 和 Niantic，以改進(jìn)我們的 PyTorch 支持。ai_edge_torch 已經(jīng)被 Shopify 團(tuán)隊(duì)用來在設(shè)備上去除產(chǎn)品圖像的背景，這個(gè)功能會(huì)出現(xiàn)在不久后發(fā)布的 Shopify 應(yīng)用中。

芯片合作伙伴和代理

我們還和 Arm、Google Tensor G3、聯(lián)發(fā)科技、高通、三星 System LSI 這些合作伙伴們一起，提供跨 CPU、GPU 和加速器的硬件支持。我們通過這些合作提高了產(chǎn)品的性能和覆蓋率，并在加速器代理上驗(yàn)證了由 PyTorch 生成的 TFLite 文件。

我們也很榮幸地和高通共同宣布新的 TensorFlow Lite 代理，現(xiàn)已開放供所有開發(fā)者使用。TFLite 代理是附加的軟件模塊，可提升在 GPU 和硬件加速器上的執(zhí)行速度。這個(gè)新的 QNN 代理支持我們?cè)?PyTorch Beta 測(cè)試集中用到的大多數(shù)模型，并提供對(duì)高通芯片的廣泛支持。通過使用高通的 DSP 和神經(jīng)處理單元，相比僅使用 CPU 和 GPU 的場(chǎng)景，能明顯地提升運(yùn)行速度 (相較 CPU 平均提升 20 倍，GPU 平均提升 5 倍)。為了方便測(cè)試，高通最近還發(fā)布了新的 AI Hub。高通 AI Hub 是一個(gè)云服務(wù)，可以讓開發(fā)者在一系列 Android 設(shè)備上對(duì) TFLite 模型進(jìn)行測(cè)試，并在使用 QNN 代理的設(shè)備上提供性能增益的可見性。

下一步

在接下來的幾個(gè)月中，我們將繼續(xù)在開放的環(huán)境中對(duì)產(chǎn)品進(jìn)行迭代，朝著 1.0 版本努力，包括提升模型覆蓋率、改進(jìn) GPU 支持，提供新的量化模式。在本系列的第二篇文章中，我們將更深入地介紹 AI Edge Torch 生成式 API，這個(gè) API 能讓開發(fā)者們?cè)谶吘壴O(shè)備中運(yùn)行自定義生成式 AI 模型，并且提供優(yōu)秀的性能表現(xiàn)。

我們要感謝所有早期體驗(yàn)用戶，正是他們提供的寶貴反饋?zhàn)屛覀兊靡约霸绨l(fā)現(xiàn)錯(cuò)誤，并確保開發(fā)者們獲得順暢的體驗(yàn)。我們還要感謝硬件合作伙伴以及 XNNPACK 生態(tài)系統(tǒng)的貢獻(xiàn)者，是他們的幫助讓我們?cè)谌绱硕鄻拥脑O(shè)備上都能獲得優(yōu)異的性能表現(xiàn)。同時(shí)，我們也要感謝廣大的 PyTorch 社區(qū)在這一路提供的指導(dǎo)和支持。