PyTorch MLP模型导出ONNX并在C++中部署的完整指南

简介将PyTorch训练好的模型部署到C服务中是许多算法落地时绕不开的一步。围绕MLP多层感知机在表格数据上的完整实现这套资源提供了一条从PyTorch导出ONNX、再到C中加载ONNX文件完成推理的清晰路线面向有模型部署需求的算法工程师与后端开发人员也适合刚接触ONNX格式的新手对照学习。资源包含PyCharm下的Python工程和VS Code gcc环境下的C工程Python侧负责模型搭建、训练和导出C侧演示如何初始化ONNX Runtime环境、读取模型并执行推理完整覆盖了MLP网络从训练到跨语言部署的关键环节。压缩包为zip格式共115个文件大小26.25MB文件类型以Python脚本、C源码、onnx模型、CMake构建配置、CSV数据集以及dll/lib/so运行库为主并包含数据文件与构建脚本方便直接复现和处理依赖。目前已有131人学习对于希望把实验模型快速迁移到业务系统的读者工程中的CMake配置和调用示例可显著减少环境配置与部署踩坑打包结构也能按需提取到自己的项目中复用。1. 为什么 MLP 训练完成后还要走一趟 ONNX一个很常见的困惑PyTorch 里 MLP 已经能跑了为什么不能直接把 .pt 文件丢给 C 程序实际上 C 侧没有 PyTorch 的 Python 运行时哪怕引入 libtorch也要在权重格式、层实现和 Tensor 生命周期上重新对齐一遍遇到自定义前向逻辑时往往很痛苦。ONNX 是一种模型交换格式它把 PyTorch 动态图固化成带权重、带算子序列的静态计算图C 端用 ONNX Runtime 加载同一份 .onnx 就能完成推理不需要重新实现 MLP 的线性层和激活函数。适合在 VS Code gcc 的工程里调用 PyTorch 训练好的模型也适合把同一个模型部署到不同平台。2. 先把 PyTorch 侧 MLP 训练到“导出即能用”不要一上来就写 torch.onnx.export。训练阶段就要为导出做准备模型结构保持简单、不使用 Dropout、BatchNorm 在导出前进入 eval 状态、只保存 state_dict 而不是整个 model。这些决定会影响 ONNX 计算图的节点数也会影响后面 C 端拿到的是不是一份干净的权重。2.1 表格数据集的读取与标准化以 CSV 表格数据为例假设有三列连续特征和一列二分类标签。先用 pandas 读入再在训练集切片上计算 mean 和 std验证集和推理阶段也必须用同一组统计量这一点在导出后特别容易被忽略。常见做法是先把 mean/std 存成 JSON 或直接写进 C 常量数组我一般不让模型内部包含标准化逻辑否则 ONNX 图里会多出若干 cast/sub/div 节点反而增加算子兼容风险。import pandas as pd import numpy as np import torch from torch.utils.data import Dataset, DataLoader # 从 CSV 读取表格数据 df pd.read_csv(tabular_dataset.csv) feature_cols [age, income, score] x df[feature_cols].values.astype(np.float32) y df[label].values.astype(np.int64) # 只在训练集切片上计算 mean/std train_num 8000 mean x[:train_num].mean(axis0, keepdimsTrue).astype(np.float32) std x[:train_num].std(axis0, keepdimsTrue).astype(np.float32) 1e-8 x (x - mean) / std class TabularDataset(Dataset): def __init__(self, x, y): self.x torch.from_numpy(x) self.y torch.from_numpy(y) def __len__(self): return len(self.x) def __getitem__(self, idx): return self.x[idx], self.y[idx] train_loader DataLoader( TabularDataset(x[:train_num], y[:train_num]), batch_size64, shuffleTrue )这里keepdimsTrue保留了二维形状后面广播到任意 batch 都方便std 1e-8是为了避免某一列零方差时除数为零。Dataset 返回的是 float32 和 int64 张量和后续torch.onnx.export的输入 dtype 保持一致。2.2 定义 MLP 网络结构MLP 的任务是把 “age / income / score” 映射到两个类别的概率分布。隐藏层用两个 32 维 Linear ReLU输出层是 2 维 logits。不要在两个隐藏层之间塞 Dropoutonnxruntime对 Dropout 不是不支持而是训练用的随机 mask 在导出后没有意义容易让人误以为是推理 bug。import torch.nn as nn class MLP(nn.Module): def __init__(self, in_dim3, hidden_dim32, out_dim2): super().__init__() self.net nn.Sequential( nn.Linear(in_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, out_dim) ) def forward(self, x): return self.net(x)这个结构导出后对应 ONNX 里的Gemm、Relu、Gemm、Relu、Gemm五个核心节点。线性层是矩阵乘法权重能被完整固化到.onnx文件里不需要在 C 端手动拿权重做 matmul。2.3 训练循环与权重保存优化器用 Adam学习率 1e-330 个 epoch 足够表格数据收敛。关键在于最后一行保存state_dict而不是torch.save(model, ...)。完整模型会带上 Python 类定义和模块路径跨机器、跨版本加载很容易出现 pickle 失败state_dict只是有序参数表导出时重新用 MLP 类构造模型再 load 即可。import torch.nn.functional as F model MLP(in_dimlen(feature_cols), hidden_dim32, out_dim2) optimizer torch.optim.Adam(model.parameters(), lr1e-3) for epoch in range(30): for xb, yb in train_loader: optimizer.zero_grad() logits model(xb) loss F.cross_entropy(logits, yb) loss.backward() optimizer.step() # 导出前必须切到 eval 状态 model.eval() torch.save(model.state_dict(), mlp.pt)训练时 BN 和 Dropout 的行为依赖全局统计量或随机状态eval 模式会把它们固化。虽然这个 MLP 里没有 BN但保留model.eval()是一种稳妥习惯导出的 ONNX 图和后续 C 推理的数值行为才会一致。3. pt 转 onnxtorch.onnx.export 的参数不是随便填的很多人转出来的 onnx 在 Python 侧能跑一到 C 就报错或输出全 0。问题通常出在动态轴设置、opset 版本或者输入 dtype 上。这一章把导出参数拆开讲。3.1 固定维度还是动态维度导出的第一步是决定输入 shape 是固定的还是动态的。ONNX 里轴可以是具体数字也可以是符号名。对于表格类 MLP 推理我通常只把 batch 维设为动态特征维固定成 3避免动态维度范围过宽导致 onnxruntime 在内存分配上走保守路径。配置方式输入 shape适用场景注意点全部固定[1, 3]单条请求、HTTP 服务最省心C 端不用关心 output shape动态 batch[batch, 3]一次处理多条样本需要dynamic_axesC 端按实际 batch 构造 tensor全动态[batch, feature]输入特征数量会变onnxruntime 支持但Gemm的 shape 推导容易出警告对小规模 MLP 来说动态 batch 和固定维度的推理性能差距几乎为 0但动态 batch 让 C 端不用为每个 batch 大小单独导出模型省一次工程参数化工作。3.2 导出代码与关键参数导出前要把模型加载到 CPU输入用torch.randn生成一个 shape 为[1, 3]的 dummy tensor。这里不要用真实 batch 很大的数据因为导出过程会真正执行一次前向dummy 输入只需要满足 shape 和 dtype 正确即可。import torch.onnx import onnx model MLP(in_dim3, hidden_dim32, out_dim2) model.load_state_dict(torch.load(mlp.pt, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, dtypetorch.float32) torch.onnx.export( model, dummy_input, mlp.onnx, export_paramsTrue, opset_version17, do_constant_foldingTrue, input_names[features], output_names[pred], dynamic_axes{ features: {0: batch}, pred: {0: batch}, }, ) onnx.checker.check_model(mlp.onnx)export_paramsTrue表示把权重也写进 ONNX 文件do_constant_foldingTrue会把 eval 状态下可以静态计算的节点折叠成常量ONNX 图更小。opset_version和 onnxruntime 版本要匹配1.14.0 的 onnxruntime 对 opset 17 支持已经稳定input_names和output_names会被 C 端原样引用后面session.Run里必须写一样的名字。如果你的模型有多个输出这里就多传几个字符串。3.3 导出后用 Python onnxruntime 做一次冒烟测试不要直接跳到 C先用 Python 侧 onnxruntime 确认导出结果和 PyTorch 原模型输出对齐。这一步能排除后面 C 工程里的环境变量、库版本等干扰。import onnxruntime as ort import numpy as np sess ort.InferenceSession(mlp.onnx, providers[CPUExecutionProvider]) test_input x[:5].astype(np.float32) pred_onnx sess.run([pred], {features: test_input})[0] with torch.no_grad(): pred_torch model(torch.from_numpy(test_input)).numpy() print(max abs diff:, np.abs(pred_onnx - pred_torch).max())如果最大绝对差异在 1e-5 量级说明导出计算图和权重都没问题。如果差异达到 1e-2先检查model.eval()是否在导出前调用再检查输入是否经过同一组 mean/std。4. 在 VS Code gcc 的 C 工程里跑 onnxruntime 推理我用的开发环境是 VS Code gcc CMake。在 Linux 下最稳妥的方式是直接把 onnxruntime 的 include 和 so 文件路径写进 CMake不要依赖系统包管理器因为你很难保证系统里的版本和训练环境一致。4.1 CMakeLists.txt 怎么指 so 路径在这个项目里onnxruntime 的库文件名是libonnxruntime.so.1.14.0头文件在同一个目录的 include 下。这里不用find_package直接指定绝对路径或相对路径构建时更直观。cmake_minimum_required(VERSION 3.16) project(mlp_onnx_runner LANGUAGES CXX) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(ONNX_RUNTIME_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/onnxruntime) include_directories(${ONNX_RUNTIME_DIR}/include) add_executable(mlp_runner main.cpp) target_link_libraries(mlp_runner ${ONNX_RUNTIME_DIR}/lib/libonnxruntime.so.1.14.0) set_target_properties(mlp_runner PROPERTIES BUILD_RPATH ${ONNX_RUNTIME_DIR}/lib )BUILD_RPATH很重要它让可执行文件运行时自动去 so 所在目录找库不用每次手动export LD_LIBRARY_PATH。如果你的环境里同时存在多个 onnxruntime 版本确认 CMake 里写死的路径是最新的否则运行时会报版本不匹配。4.2 C 端预处理与推理完整代码下面这段代码完成了“构造输入 → 建 Session → Run → 取输出”的完整闭环。注意预处理里的 mean/std 是第 2 章训练时算出来的值我在这里先用常量数组演示实际项目里建议从配置读取。#include onnxruntime_cxx_api.h #include vector #include iostream int main() { // 1. 初始化环境并加载 onnx 模型 Ort::Env env(ORT_LOGGING_LEVEL_WARNING, mlp-onnx); Ort::SessionOptions opt; opt.SetIntraOpNumThreads(1); opt.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL); Ort::Session session(env, mlp.onnx, opt); // 2. 用训练集的 mean/std 做归一化 float mean[3] {32.5f, 5600.0f, 75.0f}; float std[3] {6.2f, 1800.0f, 12.3f}; float raw[3] {34.0f, 5600.0f, 78.0f}; std::vectorfloat input(3); for (int i 0; i 3; i) { input[i] (raw[i] - mean[i]) / std[i]; } // 3. 创建 CPU 上的输入 tensor std::vectorint64_t input_shape{1, 3}; Ort::MemoryInfo mem_info Ort::MemoryInfo::CreateCpu( OrtArenaAllocator, OrtMemTypeDefault); auto input_tensor Ort::Value::CreateTensorfloat( mem_info, input.data(), input.size(), input_shape.data(), input_shape.size()); // 4. 指定输入输出节点名执行推理 const char* input_names[] {features}; const char* output_names[] {pred}; auto output_tensor session.Run( Ort::RunOptions{nullptr}, input_names, input_tensor, 1, output_names, 1); // 5. 取输出指针并打印 float* pred output_tensor.front().GetTensorMutableDatafloat(); std::cout pred[0] , pred[1] std::endl; return 0; }Ort::Env负责日志级别和线程池上下文warning 级别不会刷屏SetIntraOpNumThreads(1)在线程数不敏感的小模型上能减少调度开销。Ort::MemoryInfo::CreateCpu使用 CPU 分配器CreateTensor直接包住std::vectorfloat的内存不需要手动 copy。session.Run的两个名字数组必须和torch.onnx.export里的input_names、output_names一致否则 onnxruntime 会抛出找不到节点的异常。4.3 多 batch 调用只改 shape 和输入数据模型导出时 batch 维是动态的C 端想一次推理两条样本只要把input_shape改成{2, 3}并把 6 个 float 按行优先顺序放进input向量。输出也会是[2, 2]可以用GetShape()读取实际 shape避免把输出向量长度写死。auto shape_info output_tensor.front().GetTensorTypeAndShapeInfo(); std::vectorint64_t out_shape shape_info.GetShape(); for (auto dim : out_shape) { std::cout dim ; }这个能力来自第 3 章dynamic_axes里对 batch 维的声明。如果导出时全部维度固定这里就只能接受[1, 3]输入灵活性差很多。5. 输出对齐检查与 onnx 调优的几个实操细节关于对齐建议把 C 的输出和 Python onnxruntime 的输出对齐到同一个allclose逻辑里。C 端打印出两个类的 logitsPython 侧用相同 raw 数据跑一遍再用numpy.allclose(pred_cpp, pred_onnx, atol1e-4)判断。由于不同算子实现里的浮点累加顺序略有差异最后一位出现 1e-6 级别抖动是正常的不需要追求完全一致。5.1 一个快速定位预处理不一致的方法如果发现输出差在 0.1 以上先检查输入是否被正确归一化。在 C 端把归一化后的input向量直接打印出来和 Python 里x[:1].astype(np.float32)做逐元素比较。常见错误是把训练时的 mean/std 用成了全局统计量或者在 Python 里先做了 pandas 的fillna而 C 端没有处理缺失值导致某个特征差出一个常数。import numpy as np cpp_input np.array([1.234, 2.345, -0.567], dtypenp.float32) python_input x[0] # 训练集标准化后的第一条样本 print(np.max(np.abs(cpp_input - python_input)))如果这里输出接近 0问题就在模型文件或 onnxruntime 版本上如果输出很大说明预处理没有对齐。5.2 性能与显式图优化SetGraphOptimizationLevel(ORT_ENABLE_ALL)会启用算子融合这是 onnxruntime 默认报告里收益最大的优化项。对于 MLP 这种只有几个 Gemm 的小模型优化主要减少 CPU 调度开销真实推理时间可能从 0.1ms 降到 0.08ms感知不强但 C 工程里建议开着。还有一个值得注意的边界不要对只有一个隐藏层的 MLP 做 int8 量化。onnxruntime 1.14.0 支持部分算子量化但矩阵规模太小时反量化开销甚至会超过 FP32 前向计算。如果模型确实大于 100MB再考虑onnxruntime.quantization.quantize_dynamic量化后务必用一批真实数据对比输出均方误差。最后的工程建议把 mean、std、input_shape、onnx 路径抽到一个配置结构体里C 端从 JSON 或命令行参数加载。这样换数据集或换模型时只改配置不用重新编译也能避免把 Python 侧算出的参数散落在 main 函数各处。本文还有配套的精品资源点击获取