NVIDIA/Jetson/ 瑞芯微边缘盒子 AI 部署指南：模型转换、推理开发与测试优化

边缘计算盒子部署 AI 算法的核心是适配边缘硬件特性（算力 / 内存受限）、模型轻量化（降低资源消耗）和推理优化（提升速度），下面我会从完整流程、实操示例和关键优化点三个维度详细讲解。

一、AI 边缘计算盒子算法部署核心流程

边缘部署的整体逻辑是：训练端模型处理 → 边缘端环境适配 → 推理代码开发 → 部署测试 → 稳定运行，具体步骤如下：

1. 边缘盒子环境准备

边缘计算盒子的硬件 / 系统差异较大（如 NVIDIA Jetson、瑞芯微 RK、海思 Hi35 等），首先要完成基础环境搭建：

（1）硬件与系统确认

常见边缘盒子：

NVIDIA Jetson 系列（Nano/Xavier NX/Orin）：带 GPU，支持 CUDA/TensorRT 加速，适合中高算力场景；

瑞芯微 RK1808/RK3588：专用 NPU，适合低成本、低功耗场景；

海思 Hi3519/Hi3559：安防场景主流，支持自研 NNIE 加速。

系统：多为嵌入式 Linux（Ubuntu 20.04/18.04、BuildRoot、Yocto），优先选择盒子厂商预装的系统（自带硬件驱动）。

（2）核心依赖安装

以最常用的NVIDIA Jetson（JetPack 5.0+） 为例（预装 CUDA 11.4、CUDNN 8.6、TensorRT 8.4），补充安装必备库：

# 更新源

sudo apt update && sudo apt upgrade -y

# 安装OpenCV（处理图像/视频输入）

sudo apt install libopencv-dev python3-opencv -y

# 安装ONNX Runtime（模型转换中间件）

pip3 install onnxruntime-gpu

# 安装PyTorch（模型导出）

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu114

2. 模型预处理（边缘部署核心环节）

训练好的模型（如 PyTorch/TensorFlow）不能直接部署，需先做轻量化 + 格式转换：

（1）模型导出（转 ONNX 通用格式）

将 PyTorch 模型导出为 ONNX（跨框架、跨平台的中间格式），示例如下（以 ResNet18 图像分类模型为例）：

import torch

import torchvision.models as models

# 1. 加载预训练模型（替换为你的自定义模型）

model = models.resnet18(pretrained=True)

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

# 2. 定义输入维度（需与实际推理一致，如3通道、224x224）

input_shape = (1, 3, 224, 224)

dummy_input = torch.randn(input_shape).cuda()  # 边缘盒子有GPU则用cuda

# 3. 导出ONNX模型

onnx_path = "resnet18.onnx"

torch.onnx.export(

model,

dummy_input,

onnx_path,

opset_version=12,  # 适配TensorRT的opset版本

input_names=["input"],

output_names=["output"],

dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}  # 支持动态批量

)

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

（2）模型轻量化

针对边缘算力受限，核心优化方式：

量化：将 FP32（单精度）转为 FP16（半精度）或 INT8（整型），减少显存占用、提升推理速度（精度损失 < 5%）；

剪枝：移除模型中冗余的神经元 / 卷积核，降低模型大小；

蒸馏：用大模型（高精度）指导小模型（轻量化）训练，兼顾精度和速度。

以 TensorRT 的 INT8 量化为例（需准备校准数据集）：

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

builder = trt.Builder(TRT_LOGGER)

network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

parser = trt.OnnxParser(network, TRT_LOGGER)

# 解析ONNX模型

with open("resnet18.onnx", "rb") as f:

parser.parse(f.read())

# 配置量化（INT8）

config = builder.create_builder_config()

config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB工作空间

config.set_flag(trt.BuilderFlag.INT8)

# 加载校准数据集（需自定义Calibrator类，此处简化）

# config.int8_calibrator = MyInt8Calibrator(calib_data_path, batch_size=8)

# 构建TensorRT引擎（边缘端专用格式，后缀.engine）

serialized_engine = builder.build_serialized_network(network, config)

with open("resnet18_int8.engine", "wb") as f:

f.write(serialized_engine)

print("INT8量化后的TensorRT引擎已生成")

3. 边缘端推理代码开发

编写适配边缘盒子的推理脚本，核心逻辑：输入处理 → 模型推理 → 输出后处理，示例如下（TensorRT 推理）：

import cv2

import tensorrt as trt

import pycuda.driver as cuda

import pycuda.autoinit  # 自动初始化CUDA

import numpy as np

# 1. 加载TensorRT引擎

def load_engine(engine_path):

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:

return runtime.deserialize_cuda_engine(f.read())

# 2. 推理上下文初始化

engine = load_engine("resnet18_int8.engine")

context = engine.create_execution_context()

# 设置输入维度（与导出时一致）

context.set_binding_shape(0, (1, 3, 224, 224))

# 3. 分配CUDA内存

host_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)

host_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)

device_input = cuda.mem_alloc(host_input.nbytes)

device_output = cuda.mem_alloc(host_output.nbytes)

bindings = [int(device_input), int(device_output)]

# 4. 图像预处理（边缘端输入通常为摄像头/本地视频）

def preprocess(image_path):

img = cv2.imread(image_path)

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

img = cv2.resize(img, (224, 224))

img = img.transpose((2, 0, 1))  # HWC→CHW

img = img / 255.0  # 归一化

img = img.astype(np.float32).ravel()  # 展平

return img

# 5. 执行推理

def infer(image_path):

# 预处理

input_data = preprocess(image_path)

np.copyto(host_input, input_data)

# 数据拷贝到GPU

cuda.memcpy_htod(device_input, host_input)

# 推理

context.execute_async_v2(bindings=bindings, stream_handle=pycuda.autoinit.stream.handle)

pycuda.autoinit.stream.synchronize()

# 结果拷贝回CPU

cuda.memcpy_dtoh(host_output, device_output)

# 后处理（获取分类结果）

pred = np.argmax(host_output)

return pred

# 测试推理

if __name__ == "__main__":

result = infer("test.jpg")

print(f"推理结果（类别索引）: {result}")

# 可添加类别映射（如imagenet标签），输出具体类别名称

4. 部署与测试

（1）文件传输到边缘盒子

用scp命令将模型文件（.engine）和推理代码传到边缘盒子：

# 本地终端执行，替换为边缘盒子的IP和路径

scp resnet18_int8.engine infer.py ubuntu@192.168.1.100:/home/ubuntu/

（2）运行测试

在边缘盒子终端执行：

# 安装依赖（若未装）

pip3 install pycuda opencv-python numpy

# 运行推理脚本

python3 infer.py

（3）性能优化

查看推理速度：添加计时代码，目标是延迟 < 50ms（边缘端实时要求）；

调整批量大小：边缘端通常用 batch_size=1（实时推理）；

关闭无关进程：释放 CPU/GPU 资源，提升推理稳定性。

5. 稳定运行配置

边缘盒子需 7×24 小时运行，需配置开机自启 + 异常监控：

（1）创建 systemd 服务

编辑服务文件/etc/systemd/system/ai-infer.service：

[Unit]

Description=AI Edge Inference Service

After=network.target

[Service]

Type=simple

User=ubuntu

ExecStart=/usr/bin/python3 /home/ubuntu/infer.py

Restart=always  # 异常自动重启

RestartSec=5    # 重启间隔5秒

StandardOutput=append:/home/ubuntu/infer.log  # 日志输出

StandardError=append:/home/ubuntu/infer_error.log

[Install]

WantedBy=multi-user.target

（2）启动并设置开机自启

sudo systemctl daemon-reload

sudo systemctl start ai-infer

sudo systemctl enable ai-infer

# 查看运行状态

sudo systemctl status ai-infer

二、关键注意事项

硬件适配：不同边缘盒子的加速框架不同（NVIDIA 用 TensorRT、瑞芯微用 RKNN、海思用 NNIE），模型转换需对应厂商工具；

输入源适配：边缘端输入多为摄像头（USB/CSI），需用 OpenCV 的cv2.VideoCapture替代本地图片读取；

功耗控制：边缘盒子多为无风扇设计，需限制 CPU/GPU 频率，避免过热；

模型版本兼容：ONNX/opset 版本需与 TensorRT 版本匹配（如 TensorRT 8.4 适配 opset 12-14）。

总结

模型轻量化是核心：边缘盒子算力有限，必须通过量化 / 剪枝降低模型资源消耗，优先选择 INT8 量化（速度提升最明显）；

环境适配要精准：不同硬件（NVIDIA / 瑞芯微 / 海思）对应不同的推理框架和驱动，需匹配厂商提供的 SDK；

稳定性是落地关键：通过 systemd 配置开机自启和异常重启，结合日志监控确保 7×24 小时运行。