《Born》第19章:ONNX 导入——打通 PyTorch 与 Go 的模型桥梁
ONNX(Open Neural Network Exchange)是模型交换的事实标准。Born 的 onnx 包让你把 PyTorch 训练的模型直接加载到 Go 中运行。
---
PyTorch → ONNX → Born 完整流程
# Python 端:导出 ONNX
import torch
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx",
input_names=["input"], output_names=["output"],
opset_version=17)
// Go 端:加载并推理
backend := cpu.New()
model, err := onnx.Load("model.onnx", backend)
input := tensor.FromSlice(data, tensor.Shape{1, 3, 224, 224}, backend)
outputs, err := model.Forward(input.Raw())
---
支持的算子
Born 支持 49 个 ONNX 算子(OpSet 14-18):
| 类别 | 算子 |
|---|---|
| 数学 | Add, Sub, Mul, Div, Pow, Sqrt, Exp, Log |
| 矩阵 | MatMul, Gemm, Transpose |
| 激活 | ReLU, Sigmoid, Tanh, Softmax, Gelu, Silu |
| 形状 | Reshape, Squeeze, Unsqueeze, Concat, Split, Slice |
| 卷积 | Conv, ConvTranspose, MaxPool, AveragePool |
| 归一化 | BatchNormalization, LayerNormalization, InstanceNormalization |
不支持的算子怎么办?
1. ONNX Simplifier:onnxsim model.onnx model_simple.onnx
2. OpSet 降级:用旧版本算子语义
3. 自定义算子注册:
onnx.RegisterOp("MyCustomOp", func(inputs []tensor.Tensor) ([]tensor.Tensor, error) {
// 你的实现
})
---
📘 《Born》连载技术书,第 19/22 章。
💬 讨论回复 (0)
推荐
🌟 智谱 GLM-5 已上线
我正在智谱大模型开放平台 BigModel.cn 上打造 AI 应用,智谱新一代旗舰模型 GLM-5 已上线,在推理、代码、智能体综合能力达到开源模型 SOTA 水平。
🎁 领取 2000万 Tokens