跨平台 Go 语言 ONNX Runtime 封装库

简介

本库旨在为 Go 语言代码提供加载和执行神经网络的接口，同时尽可能保持简单易用。

使用本库的一些示例应用可以在 onnxruntime_go_examples 仓库 中找到。

onnxruntime 库提供了加载和执行 ONNX 格式神经网络的方法，但该库主要支持 C 和 C++ API。目前已有多个针对 onnxruntime 库编写的 Go 语言封装库，然而据我所知，这些现有封装库均不支持 Windows 平台。其原因在于 Microsoft 的 onnxruntime 库假设用户将在 Windows 系统上使用 MSVC 编译器，而 Windows 上的 CGo 则要求使用 Mingw。

本封装库通过手动加载 onnxruntime 共享库来绕过这些问题，从而不再依赖于 onnxruntime 源代码，仅需头文件即可。当然，这种方法同样适用于非 Windows 系统。

此外，本库利用 Go 语言最近引入的泛型特性，支持多种张量数据类型；可参阅 NewTensor 或 NewEmptyTensor 函数。

经过测试并确认，包括 TensorRT、CUDA 和 CoreML 在内的多种加速执行提供者均可与 onnxruntime_go 配合使用。本 README 的“要求”部分（见下文）提供了更多详细信息。

关于 onnxruntime 库版本的说明

截至撰写本文时，本库使用的是 onnxruntime C API 头文件的 1.24.1 版本。因此，它很可能也只能与 onnxruntime 共享库的 1.24.1 版本配合使用。如果您需要使用其他版本，或者当我未能及时更新此仓库时，更新或更改 onnxruntime 版本应该相对容易：

1. 将 onnxruntime_c_api.h 和 onnxruntime_ep_c_api.h 文件替换为您希望使用的 onnxruntime 版本对应的文件。
2. 将 test_data/onnxruntime.dll（或 test_data/onnxruntime*.so、test_data/onnxruntime*.dylib）文件替换为您希望使用的 onnxruntime 版本对应的文件。
3. （如果您关心 DirectML 支持）请验证 onnxruntime_wrapper.c 中 DummyOrtDMLAPI 结构体中的条目顺序是否与官方仓库中 dml_provider_factory.h 头文件中 OrtDmlApi 结构体中的顺序一致。更多信息请参阅 onnxruntime_wrapper.c 中对该结构体的注释。

请注意，C API 头文件和共享库文件均可从 官方仓库 的发布页面下载。下载您希望使用的版本的压缩包并解压即可。头文件位于 “include” 子目录中，共享库则位于 “lib” 子目录中。（在 Linux 系统上，您需要带有版本号后缀的 .so 文件，例如 libonnxruntime.so.1.24.1，而不是仅仅作为符号链接的 libonnxruntime.so。）压缩包中还包含其他 C++ 头文件、调试符号等文件，但您通常只需要单个 onnxruntime 共享库以及两个 _c_api.h 头文件即可。（例外情况是如果您希望启用 GPU 支持，可能还需要其他共享库文件，如 execution_providers_cuda.dll 和 execution_providers_shared.dll（或其在 Linux 或 OSX 上的对应文件）。）

要求

要使用本库，您需要一个支持 cgo 的 Go 版本。此外，您还需要根据您的操作系统和架构获取正确版本的 onnxruntime 共享库或 DLL。在初始化 onnxruntime_go 之前，您需要提供该共享库的路径。请参阅以下示例的前几行代码（即 ort.SetSharedLibraryPath(...)）。

如果您希望使用 CUDA，您不仅需要使用支持 CUDA 的 onnxruntime 共享库版本，还需确保所使用的 CUDA 版本与您的 onnxruntime 库底层版本兼容。例如，onnxruntime 1.23.2 版本仅支持 CUDA 12.x 版本。有关更具体的信息，请参阅 onnxruntime CUDA 支持文档。

与 CUDA 类似，其他执行提供者也有各自独立的使用要求。这些要求过于繁多，无法在此 README 中一一列出。请务必先在 Python 脚本中成功使用您选择的执行提供者，再在此处提出相关问题。

使用示例

完整的文档可在 pkg.go.dev 上查阅。

此外，几个包含必要网络和数据的命令行示例应用程序也可在 onnxruntime_go_examples 仓库 中找到。

以下示例展示了如何使用本库加载并运行一个接受单个输入张量、输出单个张量的 ONNX 网络，这两个张量均包含 32 位浮点数值。请注意，此处省略了错误处理；每个函数都会返回一个 err 值，若发生失败则该值将非空。

import (
    "fmt"
    ort "github.com/yalue/onnxruntime_go"
    "os"
)

func main() {
    // 此行 _可能_ 是可选的；默认情况下，库会在 Windows 上尝试加载 "onnxruntime.dll"，而在其他系统上尝试加载 "onnxruntime.so"。为保证稳定性，建议程序始终显式设置此路径。
    ort.SetSharedLibraryPath("path/to/onnxruntime.so")

    err := ort.InitializeEnvironment()
    if err != nil {
        panic(err)
    }
    defer ort.DestroyEnvironment()

    // 为了略微提升性能并方便重复使用现有张量，本库要求用户在创建会话之前先创建所有输入和输出张量。如果这对您的应用场景不太合适，可以参考文档中的 DynamicAdvancedSession 类型，它允许在调用 Run() 时而非初始化会话时指定输入和输出张量。
    inputData := []float32{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9}
    inputShape := ort.NewShape(2, 5)
    inputTensor, err := ort.NewTensor(inputShape, inputData)
    defer inputTensor.Destroy()
    // 该假设网络将 2x5 的输入映射为 2x3x4 的输出。
    outputShape := ort.NewShape(2, 3, 4)
    outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
    defer outputTensor.Destroy()

session, err := ort.NewAdvancedSession("path/to/network.onnx",
        []string{"Input 1 Name"}, []string{"Output 1 Name"},
        []ort.Value{inputTensor}, []ort.Value{outputTensor}, nil)
    defer session.Destroy()

    // 调用 Run() 将运行网络，读取输入张量的当前内容并修改输出张量的内容。
    err = session.Run()

    // 获取输出张量数据的切片视图。
    outputData := outputTensor.GetData()

    // 如果你想使用不同的输入运行网络，只需修改输入张量的数据（可通过 inputTensor.GetData() 获取）
    // 然后再次调用 Run() 即可。

    // ...
}

已弃用的 API

类型化的 Session[t]: 该库的旧版本使用了一个带类型的 Session[T] 结构体来管理会话。回过头来看，将类型参数与会话关联起来是不必要的，因此引入了 AdvancedSession 类型及其相关 API 来纠正这一问题。出于向后兼容性的考虑，旧的类型化 Session[T] 和 DynamicSession[T] 类型仍然保留，并且短期内不太可能被移除。不过，它们现在会在内部将功能委托给 AdvancedSession。建议新代码始终直接使用 AdvancedSession。

ONNX Runtime 的训练 API: 训练 API 自 ONNX Runtime 1.20 版本起已被弃用。为了不再维护针对已弃用 API 的封装，onnxruntime_go 已将训练 API 的封装函数替换为返回错误的存根函数。如果用户仍需使用训练 API，则需要使用较旧的版本。例如，以下版本应与训练功能兼容：

• onnxruntime_go 的 v1.12.1 版本，
• ONNX Runtime 的 1.19.2 版本。

运行测试及测试的系统兼容性

请导航到此目录并运行 go test -v，或者可选地运行 go test -v -bench=.。所有测试都应通过；与 CUDA 或其他加速器支持相关的测试将在不支持这些功能的系统或 ONNX Runtime 构建中被跳过。

目前，该仓库包含适用于几种系统的 ONNX Runtime 共享库副本，包括 AMD64 Windows、ARM64 Linux 和 ARM64 Darwin。这些库应该能够让这些系统上的测试顺利通过，而无需用户在克隆仓库之外再复制额外的库文件。然而，未来随着对更多或更少系统的支持增加或减少，这种情况可能会发生变化。

您可能出于以下原因希望使用不同版本的 ONNX Runtime 共享库：

1. 仓库中包含的共享库副本并不支持 CUDA 或其他加速执行提供者，因此如果您使用的是默认库，与 CUDA 相关的测试将始终被跳过。

2. 许多系统，包括 AMD64 和 i386 Linux 以及 x86 macOS，目前根本未在 test_data/ 中包含共享库。（我希望通过尽量减少共享库的数量来保持该目录及整个仓库的体积较小。）

如果上述情况或其他原因适用于您，测试代码会在尝试从 test_data/ 加载库之前检查 ONNXRUNTIME_SHARED_LIBRARY_PATH 环境变量。因此，如果您正在使用这些系统之一，或者希望运行与加速器相关的测试，您应当将该环境变量设置为 ONNX Runtime 共享库的路径。之后，运行 go test -v 应当能够成功通过所有测试。

onnxruntime_go 快速上手指南

环境准备

系统要求

• Go 版本：需要支持 cgo 的 Go 版本（建议使用 Go 1.18+ 以利用泛型特性）。
• 操作系统：支持 Windows、Linux 和 macOS。
  • Windows 需注意：本库通过手动加载共享库规避了 MSVC 与 Mingw 的编译器冲突问题，因此可在标准 Go 环境下运行。
• 架构支持：AMD64、ARM64 等主流架构。

前置依赖

1. ONNX Runtime 共享库：

   • 需下载与代码中头文件版本匹配（当前为 1.24.1）的共享库文件。
   • Windows: onnxruntime.dll
   • Linux: libonnxruntime.so.1.24.1 (注意需带版本号的文件，而非符号链接)
   • macOS: libonnxruntime.1.24.1.dylib
   • 下载地址：Microsoft ONNX Runtime Releases。下载对应版本的压缩包，从 lib 目录获取共享库文件。

2. GPU 加速支持（可选）：

   • 若需使用 CUDA、TensorRT 或 CoreML，请下载包含相应 Execution Provider 支持的 ONNX Runtime 版本，并确保系统已安装对应的驱动（如 CUDA Toolkit）。

安装步骤

1. 初始化 Go 模块并安装库：

   go get github.com/yalue/onnxruntime_go
   

2. 配置共享库路径： 在运行程序前，需确保程序能找到 ONNX Runtime 共享库。有两种方式：

   • 方式一：代码中指定（推荐） 在代码初始化环境中显式设置路径（见下方“基本使用”示例）。

   • 方式二：环境变量指定（用于测试） 如果运行测试或希望全局指定，可设置环境变量：

     # Linux/macOS
     export ONNXRUNTIME_SHARED_LIBRARY_PATH=/path/to/libonnxruntime.so.1.24.1
     
     # Windows (PowerShell)
     $env:ONNXRUNTIME_SHARED_LIBRARY_PATH="C:\path\to\onnxruntime.dll"
     

基本使用

以下示例展示如何加载一个 ONNX 模型并执行推理。假设模型接受一个 2x5 的浮点输入，输出一个 2x3x4 的浮点结果。

package main

import (
    "fmt"
    ort "github.com/yalue/onnxruntime_go"
    "os"
)

func main() {
    // 1. 设置共享库路径 (为了稳定性，建议始终显式设置)
    // Windows 下通常为 "onnxruntime.dll"，其他系统为 "onnxruntime.so"
    ort.SetSharedLibraryPath("path/to/onnxruntime.so")

    // 2. 初始化环境
    err := ort.InitializeEnvironment()
    if err != nil {
        panic(err)
    }
    defer ort.DestroyEnvironment()

    // 3. 准备输入数据 (2x5 的浮点数组)
    inputData := []float32{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9}
    inputShape := ort.NewShape(2, 5)
    inputTensor, err := ort.NewTensor(inputShape, inputData)
    if err != nil {
        panic(err)
    }
    defer inputTensor.Destroy()

    // 4. 创建空输出张量 (2x3x4)
    outputShape := ort.NewShape(2, 3, 4)
    outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
    if err != nil {
        panic(err)
    }
    defer outputTensor.Destroy()

    // 5. 创建会话 (需替换为实际的模型路径及输入/输出节点名称)
    session, err := ort.NewAdvancedSession("path/to/network.onnx",
        []string{"Input 1 Name"}, []string{"Output 1 Name"},
        []ort.Value{inputTensor}, []ort.Value{outputTensor}, nil)
    if err != nil {
        panic(err)
    }
    defer session.Destroy()

    // 6. 执行推理
    err = session.Run()
    if err != nil {
        panic(err)
    }

    // 7. 获取输出结果
    outputData := outputTensor.GetData()
    fmt.Println("Output data:", outputData)
}

关键点说明

• 张量管理：输入和输出张量需在创建 Session 前准备好。使用 defer tensor.Destroy() 及时释放内存。
• 泛型支持：使用 NewEmptyTensor[T] 可灵活指定数据类型（如 float32, int64 等）。
• 重复推理：若需对同一模型进行多次推理，只需修改 inputTensor.GetData() 中的数据并再次调用 session.Run() 即可，无需重建 Session。