当前位置：首页 > news >正文

【RK3588实战】从PyTorch到嵌入式部署：一个图像分类模型的完整落地之旅

news 2026/4/17 7:39:19

1. 从零开始：PyTorch模型训练实战

第一次接触RK3588部署的朋友可能会疑惑：为什么非要先训练一个PyTorch模型？直接拿现成的模型不行吗？其实这里有个关键认知——模型部署的本质是将训练好的算法适配到特定硬件。就像你要把家具搬进新家，总得先有家具吧？下面我就用CIFAR-10数据集为例，手把手带你打造自己的"家具"。

先说说环境配置的坑。很多教程一上来就让你装CUDA，但实际测试发现，PyTorch 1.8+版本对CUDA 11.6的支持最稳定。我的开发环境是这样的：

Ubuntu 20.04 LTS
Python 3.8.10
PyTorch 1.12.1 + CUDA 11.6
torchvision 0.13.1

模型结构设计有个小技巧：先做减法再做加法。新手常犯的错误是盲目堆叠网络层，结果训练时显存爆炸。我设计的这个CNN结构经过实测，在RK3588上既能保证精度又能流畅运行：

class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.fc1 = nn.Linear(64 * 4 * 4, 500) self.fc2 = nn.Linear(500, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2(x), 2)) x = F.relu(F.max_pool2d(self.conv3(x), 2)) x = x.view(-1, 64 * 4 * 4) x = F.relu(self.fc1(x)) return self.fc2(x)

训练过程中的三个关键点：

数据预处理一致性： transforms.Normalize的参数(0.5,0.5,0.5)后面会直接影响RKNN的配置
设备切换陷阱：记得用.to(device)把模型和数据都放到GPU上
模型保存方式：建议同时保存state_dict和整个模型，后者方便后续调试

2. PyTorch到ONNX：模型转换的暗礁

转换ONNX时踩过的坑，可能比训练模型还多。最常见的就是版本兼容性问题。有次我用PyTorch 1.10导出的ONNX，到RKNN工具链里直接报错，最后发现是opset_version设成了12，改成11立即解决。

导出时的核心参数解析：

torch.onnx.export( model, torch.randn(1, 3, 32, 32, device=device), # 注意batch_size=1 "model.onnx", export_params=True, opset_version=11, # RKNN目前最佳兼容版本 do_constant_folding=True, input_names=["input"], output_names=["output"], dynamic_axes=None # RK3588需要静态shape )

验证ONNX模型时有个容易忽略的点：PC端推理要与训练预处理严格一致。有次测试准确率暴跌，查了半天发现是Resize时用了不同的插值方法。推荐用这个验证脚本：

def validate_onnx(): ort_session = ort.InferenceSession("model.onnx") # 用测试集数据验证 correct = 0 for data, target in test_loader: outputs = ort_session.run( None, {"input": data.numpy()} ) pred = outputs[0].argmax(1) correct += (pred == target.numpy()).sum() print(f"ONNX验证准确率: {100 * correct / len(test_loader.dataset):.1f}%")

3. ONNX到RKNN：嵌入式适配的关键一跃

转换RKNN模型时，90%的问题都出在config配置上。特别是mean_values和std_values，这里有个数值映射的玄机：

rknn.config( mean_values=[[127.5, 127.5, 127.5]], # (0.5,0.5,0.5)映射到0-255范围 std_values=[[127.5, 127.5, 127.5]], # 同上 target_platform="rk3588" )

量化是个双刃剑。虽然能提升推理速度，但处理不当会导致精度悬崖式下跌。我的经验是：

准备至少500张校准图片
数据集要覆盖所有类别
量化后必须做精度验证

转换流程的完整代码示例：

rknn = RKNN(verbose=True) ret = rknn.load_onnx(model="model.onnx") ret = rknn.build(do_quantization=True, dataset="calib.txt") ret = rknn.export_rknn("model.rknn") rknn.release()

4. RK3588部署实战：Python与C++双方案

Python部署最适合快速验证，但要注意RKNNLite的内存管理：

rknn_lite = RKNNLite() rknn_lite.load_rknn("model.rknn") rknn_lite.init_runtime() # 必须确保输入数据是NHWC格式 outputs = rknn_lite.inference(inputs=[img_np]) rknn_lite.release() # 这个容易忘！

C++部署的性能更优，但要注意：

输入数据排布必须是NHWC
OpenCV默认读入是BGR需要转换
内存管理要手动释放

关键代码片段：

cv::Mat img = cv::imread("test.jpg"); cv::cvtColor(img, img, cv::COLOR_BGR2RGB); cv::resize(img, img, cv::Size(32, 32)); float* input_data = (float*)malloc(32*32*3*sizeof(float)); // 手动填充数据到input_data rknn_inputs_set(ctx, 1, inputs); rknn_run(ctx, nullptr); rknn_outputs_get(ctx, 1, outputs, nullptr);