ExecuTorch 实战：在 Arm CPU 与 NPU 上跑起端侧 AI 推理

# 创建工作目录并安装 ExecuTorch
python -m venv executorch_env
source executorch_env/bin/activate

pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install executorch

# 安装 Arm XNNPACK 后端依赖
pip install executorch-backends-xnnpack

# 安装 Ethos-U 后端（需要 aarch64 交叉编译工具链）
# 注意：Ethos-U 编译步骤依赖 Arm 的 Vela 工具，先安装它
pip install ethos-u-vela

import torch
from executorch.exir import to_edge

# 用一个轻量 MobileNetV2 示例；实际可替换为你自己的模型
model = torch.hub.load("pytorch/vision:v0.16", "mobilenet_v2", pretrained=True)
model.eval()

# 标准导出：捕获静态计算图
sample_input = (torch.randn(1, 3, 224, 224),)
exported_program = torch.export.export(model, sample_input)

# 转为 Edge Dialect IR
edge_program = to_edge(exported_program)

# 保存中间表示，后续可分别编译给不同后端
edge_program.save("mobilenet_v2_edge.pte")
print("Edge Dialect 导出完成")

from executorch.exir import to_edge
from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import XnnpackQuantizer
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e

# 重新加载导出结果（独立脚本场景）
exported_program = torch.export.export(model, sample_input)

# XNNPACK 量化：8-bit 权重 + 激活量化
quantizer = XnnpackQuantizer()
prepared_program = prepare_pt2e(exported_program, quantizer)

# 用校准数据跑几轮，收集量化统计
for _ in range(3):
    prepared_program.model()(torch.randn(1, 3, 224, 224))

converted_program = convert_pt2e(prepared_program)

# 转 Edge + 分区给 XNNPACK
edge_program = to_edge(converted_program)
edge_program = edge_program.to_executorch(XnnpackPartitioner)

# 生成最终 PTE 文件，拷到设备上即可运行
with open("mobilenet_v2_xnnpack.pte", "wb") as f:
    f.write(edge_program.buffer)
print("XNNPACK 编译完成，输出 mobilenet_v2_xnnpack.pte")

# 交叉编译 runtime（以 aarch64 为例）
git clone https://github.com/pytorch/executorch.git
cd executorch
mkdir build && cd build

cmake .. \
  -DCMAKE_TOOLCHAIN_FILE=../cmake/aarch64-linux-gnu.toolchain.cmake \
  -DEXECUTORCH_BUILD_XNNPACK=ON \
  -DEXECUTORCH_BUILD_ETHOS_U=ON

cmake --build . -j$(nproc)

# 把生成的 runtime 二进制和 .pte 文件拷到目标板卡
scp libexecutorch.so mobilenet_v2_xnnpack.pte user@arm-board:/tmp/

#include <executorch/runtime/executor/program.h>
#include <executorch/runtime/executor/method.h>

using namespace executorch::runtime;

int main() {
    // 加载 PTE 文件
    Program program = Program::load("/tmp/mobilenet_v2_xnnpack.pte");
    Method method = program.loadMethod("forward");

    // 准备输入 tensor（EValue 是 ExecuTorch 的统一数据容器）
    float input_data[1 * 3 * 224 * 224];
    EValue input = EValue(torch::from_blob(
        input_data, {1, 3, 224, 224}, torch::kFloat32));

    // 执行推理
    std::vector<EValue> outputs;
    method.execute({input}, &outputs);

    // outputs[0] 即分类 logits
    printf("推理完成，输出 tensor 大小: %ld\n",
           outputs[0].toTensor().numel());
    return 0;
}