使用onnx-go实现ONNX接口的Golang插件库

ONNX (Open Neural Network Exchange) 是一种用于表示深度学习模型的开放格式，onnx-go是一个Golang实现的ONNX运行时库，允许在Go环境中加载和运行ONNX模型。

安装onnx-go

首先安装onnx-go库：

go get github.com/owulveryck/onnx-go

基本使用示例

1. 加载ONNX模型

package main

import (
	"fmt"
	"log"
	"os"

	"github.com/owulveryck/onnx-go"
	"gorgonia.org/tensor"
)

func main() {
	// 1. 加载ONNX模型文件
	model := onnx.NewModel(onnx.NewGraph())
	
	file, err := os.Open("model.onnx")
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	
	err = model.UnmarshalBinary(file)
	if err != nil {
		log.Fatal(err)
	}

	// 2. 准备输入数据
	inputT := tensor.New(
		tensor.WithShape(1, 3, 224, 224), // 示例输入形状
		tensor.WithBacking([]float32{...}), // 填入实际输入数据
	)

	// 3. 创建输入输出映射
	input := make(map[string]interface{})
	output := make(map[string]interface{})
	
	// 获取输入节点名称(通常需要查看模型文档)
	inputName := model.Graph().Input[0].Name
	input[inputName] = inputT
	
	// 4. 运行模型
	err = model.Run(input, output)
	if err != nil {
		log.Fatal(err)
	}

	// 5. 获取输出
	for name, value := range output {
		fmt.Printf("Output %v: %v\n", name, value)
	}
}

2. 使用后端加速

onnx-go支持不同的后端计算引擎，例如Gorgonia：

import (
	"github.com/owulveryck/onnx-go/backend/x/gorgonnx"
)

func main() {
	// 创建带有Gorgonia后端的模型
	backend := gorgonnx.NewGraph()
	model := onnx.NewModel(backend)
	
	// 其余代码与前面相同...
}

高级用法

1. 处理多输入输出

// 准备多个输入
input1 := tensor.New(tensor.WithShape(1, 3), tensor.WithBacking([]float32{1, 2, 3}))
input2 := tensor.New(tensor.WithShape(1, 2), tensor.WithBacking([]float32{4, 5}))

inputs := map[string]interface{}{
	"input1": input1,
	"input2": input2,
}

outputs := make(map[string]interface{})

err := model.Run(inputs, outputs)
if err != nil {
	log.Fatal(err)
}

for name, value := range outputs {
	fmt.Printf("%s: %v\n", name, value)
}

2. 获取模型元信息

// 获取模型输入输出信息
fmt.Println("Model inputs:")
for _, input := range model.Graph().Input {
	fmt.Printf("  Name: %v\n", input.Name)
	fmt.Printf("  Type: %v\n", input.Type)
}

fmt.Println("Model outputs:")
for _, output := range model.Graph().Output {
	fmt.Printf("  Name: %v\n", output.Name)
	fmt.Printf("  Type: %v\n", output.Type)
}

注意事项

1. 性能考虑：Golang不是深度学习的主流语言，onnx-go的性能可能不如Python实现
2. 操作符支持：不是所有ONNX操作符都被完全支持，使用前需检查兼容性
3. 模型转换：确保模型已正确转换为ONNX格式
4. 输入输出形状：需要与模型期望的形状完全匹配

完整示例：图像分类

package main

import (
	"fmt"
	"log"
	"os"

	"github.com/owulveryck/onnx-go"
	"github.com/owulveryck/onnx-go/backend/x/gorgonnx"
	"gorgonia.org/tensor"
)

func main() {
	// 1. 初始化模型
	backend := gorgonnx.NewGraph()
	model := onnx.NewModel(backend)

	// 2. 加载ONNX模型
	file, err := os.Open("resnet18.onnx")
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()
	
	err = model.UnmarshalBinary(file)
	if err != nil {
		log.Fatal(err)
	}

	// 3. 准备输入数据 (假设是1x3x224x224的图片)
	inputData := make([]float32, 1*3*224*224)
	// 这里应该填充实际的图像数据，通常需要预处理(归一化等)
	
	inputT := tensor.New(
		tensor.WithShape(1, 3, 224, 224),
		tensor.WithBacking(inputData),
	)

	// 4. 运行推理
	input := map[string]interface{}{"input.1": inputT}
	output := map[string]interface{}{}
	
	err = model.Run(input, output)
	if err != nil {
		log.Fatal(err)
	}

	// 5. 处理输出 (假设输出是1x1000的类别概率)
	outputTensor := output["output"].(*tensor.Dense)
	probabilities := outputTensor.Data().([]float32)
	
	// 找到最高概率的类别
	maxIndex, maxProb := 0, float32(0)
	for i, prob := range probabilities {
		if prob > maxProb {
			maxIndex, maxProb = i, prob
		}
	}
	
	fmt.Printf("Predicted class: %d with probability: %.2f%%\n", 
		maxIndex, maxProb*100)
}

onnx-go为Golang开发者提供了在Go生态中使用ONNX模型的途径，虽然功能可能不如Python生态完善，但对于需要集成深度学习模型到Go应用中的场景是一个有价值的工具。