BM1684X AI芯片边缘计算盒子：PyTorch 和 TensorFlow 模型部署详解

算能 BM1684X作为一款面向边缘计算场景的AI芯片，确实支持PyTorch/TensorFlow框架，但并非直接原生运行，而是需要通过算能提供的工具链进行模型转换和部署。下面我会详细讲解具体的支持方式、部署流程和实操代码，帮助你快速上手。

一、核心支持逻辑

BM1684X 有自己的指令集和计算架构，无法直接运行PyTorch/TensorFlow的原生模型（如.pth/.pb文件），需要通过算能的Sophon SDK 将框架模型转换为芯片可执行的bmodel格式，整体流程如下：

二、环境准备

1. 硬件与系统要求

硬件：BM1684X 边缘计算盒子（如万物纵横DA320S/DA380SF）

系统：Ubuntu 18.04/20.04（推荐）

依赖：Sophon SDK（算能官网下载对应版本）

2. 安装Sophon SDK

# 解压SDK包（以sophon-soc-sdk_v2.7.0为例）

tar -xvf sophon-soc-sdk_v2.7.0_ubuntu18.04.tar.gz

cd sophon-soc-sdk_v2.7.0_ubuntu18.04

# 安装基础依赖

sudo ./install_deps.sh

# 安装SDK核心组件

sudo ./install.sh

三、PyTorch模型部署实操（以ResNet50为例）

1. 导出PyTorch模型为ONNX格式

首先将训练好的PyTorch模型转换为ONNX（通用中间格式），这是对接BM1684X的关键步骤：

import torch

import torchvision.models as models

# 1. 加载预训练的ResNet50模型

model = models.resnet50(pretrained=True)

model.ｅｖａｌ()  # 切换到推理模式

# 2. 构造输入示例（需与实际推理输入尺寸一致）

input_tensor = torch.randn(1, 3, 224, 224)  # batch=1, 3通道, 224x224

# 3. 导出ONNX模型

onnx_path = "resnet50.onnx"

torch.onnx.export(

    model,

    input_tensor,

    onnx_path,

    opset_version=11,  # 推荐opset 11/12，兼容Sophon SDK

    input_names=["input"],

    output_names=["output"],

    dynamic_axes=None  # BM1684X暂不推荐动态维度，固定尺寸更稳定

)

print(f"ONNX模型已导出至: {onnx_path}")

2. 将ONNX转换为BM1684X的bmodel格式

使用Sophon SDK的bmnetp工具（专门处理PyTorch/ONNX模型）进行转换：

# 转换命令（适配BM1684X）

bmnetp --model=resnet50.onnx \

       --weight=resnet50.onnx \

       --shapes=[1,3,224,224] \

       --net_name=resnet50 \

       --target=BM1684X \

       --precision=fp32 \  # 可选fp16/int8，int8需量化

       --outdir=./bmodel_output

# 转换完成后，在bmodel_output目录下会生成resnet50.bmodel文件

3. 在BM1684X上加载bmodel推理

使用Sophon SDK的bmlib和bmruntime库进行推理：

import cv2

import numpy as np

from sophon import bmruntime

# 1. 初始化BM1684X运行时

dev_id = 0  # 边缘盒的设备ID，通常为0

runtime = bmruntime.BMRuntime(dev_id)

# 2. 加载bmodel模型

bmodel_path = "./bmodel_output/resnet50.bmodel"

net_id = runtime.load_bmodel(bmodel_path)

graph_names = runtime.get_graph_names(net_id)

graph_name = graph_names[0]

# 3. 预处理输入图片（与训练时的预处理一致）

img_path = "test.jpg"

img = cv2.imread(img_path)

img = cv2.resize(img, (224, 224))  # 缩放到模型输入尺寸

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR转RGB

img = np.transpose(img, (2, 0, 1))  # HWC -> CHW

img = np.expand_dims(img, 0)  # 增加batch维度

img = img.astype(np.float32) / 255.0  # 归一化

# 4. 构造BM1684X的输入张量

input_tensor = runtime.get_input_tensor(graph_name, "input")

input_tensor.from_numpy(img)

# 5. 执行推理

output = runtime.run(graph_name, {"input": input_tensor})

# 6. 解析输出（获取分类结果）

pred = np.argmax(output[0].numpy())

print(f"推理结果类别ID: {pred}")

四、TensorFlow模型部署（补充说明）

TensorFlow模型的部署流程与PyTorch类似，核心差异在第一步：

1. 将TensorFlow模型（.ckpt/.h5）导出为PB格式或直接转换为ONNX：

import tensorflow as tf

from tensorflow.keras.applications.resnet50 import ResNet50

# 加载模型并导出为ONNX（需安装tf2onnx）

model = ResNet50(weights="imagenet")

tf.saved_model.save(model, "./resnet50_tf")

# 终端执行转换命令

# python -m tf2onnx.convert --saved-model ./resnet50_tf --output resnet50_tf.onnx

2. 后续使用bmnett工具（TensorFlow专用）转换为bmodel：

bmnett --model=resnet50_tf.onnx --shapes=[1,224,224,3] --target=BM1684X --outdir=./bmodel_output

3. 推理代码与PyTorch转换后的bmodel完全一致。

五、关键注意事项

1. 量化优化：BM1684X对int8量化模型的性能提升显著（算力利用率更高），可通过Sophon SDK的calibration_tool工具进行量化；

2. 算子兼容：部分PyTorch/TensorFlow的小众算子可能不支持，需提前查看Sophon SDK的算子支持列表，不兼容算子需手动替换；

3. 性能调优：可通过设置batch_size、使用多线程推理提升边缘盒的吞吐量。

总结

1. BM1684X 不直接运行PyTorch/TensorFlow原生模型，需通过Sophon SDK将模型转为bmodel格式；

2. 核心流程为：框架模型 → ONNX → bmodel → BM1684X推理，PyTorch用bmnetp、TensorFlow用bmnett转换；

3. 推理阶段依赖Sophon SDK的bmruntime库，输入需严格匹配模型预处理逻辑。