TVM Golang 运行时，用于深度学习部署

•


简介

TVM 是一个开放的深度学习编译器堆栈，用于将来自不同框架的各种深度学习模型编译到 CPU、GPU 或专用加速器。TVM 支持从各种前端（如 Tensorflow、Onnx、Keras、Mxnet、Darknet、CoreML 和 Caffe2）进行模型编译。TVM 编译的模块可以部署在后端，如 LLVM (Javascript 或 WASM、AMD GPU、ARM 或 X86)、NVidia GPU (CUDA)、OpenCL 和 Metal。

TVM 支持 Javascript、Java、Python、C++ 等编程语言的运行时绑定，现在也支持 Golang。凭借广泛的前端、后端和运行时绑定，TVM 使开发人员能够通过多种编程语言将来自各种框架的深度学习模型集成和部署到各种硬件上。

TVM 导入和编译过程会生成一个图 JSON、一个模块和一个参数 (params)。任何集成了 TVM 运行时的应用程序都可以加载这些编译的模块并执行推理。有关使用 TVM 进行模块导入和编译的详细教程，请访问 tutorials。

TVM 现在支持通过 Golang 部署编译的模块。Golang 应用程序可以利用这一点通过 TVM 部署深度学习模型。本博客的范围是介绍 gotvm 包、包构建过程以及使用 gotvm 加载编译模块并执行推理的示例应用程序。

包

golang 包 gotvm 构建于 TVM 的 C 运行时接口之上。此包中的 API 抽象了原生 C 类型，并提供了与 Golang 兼容的类型。包源代码可以在 gotvm 中找到。

此包利用了 golang 的接口、切片、函数闭包，并隐式处理 API 调用之间必要的转换。

image

Golang 接口，基于 TVM 运行时

如何

如下图所示，gotvm 使 golang 应用程序能够集成来自各种框架的深度学习模型，而无需费力理解每个框架相关的接口 API。开发人员可以利用 TVM 导入和编译深度学习模型并生成 TVM 工件。gotvm 包提供了 golang 友好的 API 来加载、配置、输入数据和获取输出。

image

导入、编译、集成和部署

TVM 编译深度学习模型 教程可用于编译来自 TVM 前端支持的所有框架的模型。此编译过程会生成在目标上集成和部署模型所需的工件。

API

gotvm 包提供了一些数据类型和 API 函数，用于从 golang 应用程序初始化、加载和推理。像任何其他 golang 包一样，我们只需要在此处导入 gotvm 包。

  • Module（模块）：Module API 可用于将 TVM 编译的模块加载到 TVM 运行时并访问任何函数。
  • Value（值）：Value API 提供辅助函数，用于在 golang 类型（如基本类型或切片）中设置参数或获取返回值。
  • Function（函数）：Function API 对于获取函数句柄并调用它们很有用。
  • Array（数组）：Array API 对于通过 golang 切片设置和获取 Tensor 数据很有用。
  • Context（上下文）：Context API 包含辅助函数，用于构建后端上下文句柄。

示例

下面展示了一个简单的示例，其中包含加载编译模块和执行推理的内联文档。为简单起见，此处忽略了错误处理，但在实际应用中这很重要。


package main

// Import compiled gotvm package.
import (
    "./gotvm"
)

// Some constants for TVM compiled model paths.
// modLib : Is the compiled library exported out of compilation.
// modJson : TVM graph JSON.
// modParams : Exported params out of TVM compilation process.
const (
    modLib    = "./libdeploy.so"
    modJSON   = "./deploy.json"
    modParams = "./deploy.params"
)

// main
func main() {
    // Some util API to query underlying TVM and DLPack version information.
    fmt.Printf("TVM Version   : v%v\n", gotvm.TVMVersion)
    fmt.Printf("DLPACK Version: v%v\n\n", gotvm.DLPackVersion)

    // Import tvm module (so).
    modp, _ := gotvm.LoadModuleFromFile(modLib)

    // Load module on tvm runtime - call tvm.graph_runtime.create
    // with module and graph JSON.
    bytes, _ := ioutil.ReadFile(modJSON)
    jsonStr := string(bytes)
    funp, _ := gotvm.GetGlobalFunction("tvm.graph_runtime.create")
    graphrt, _ := funp.Invoke(jsonStr, modp, (int64)(gotvm.KDLCPU), (int64)(0))
    graphmod := graphrt.AsModule()


    // Allocate input & output arrays and fill some data for input.
    tshapeIn  := []int64{1, 224, 224, 3}
    tshapeOut := []int64{1, 1001}
    inX, _ := gotvm.Empty(tshapeIn, "float32", gotvm.CPU(0))
    out, _ := gotvm.Empty(tshapeOut)
    inSlice := make([]float32, (244 * 244 * 3))
    rand.Seed(10)
    rand.Shuffle(len(inSlice), func(i, j int) {inSlice[i],
                                               inSlice[j] = rand.Float32(),
                                               rand.Float32() })
    inX.CopyFrom(inSlice)

    // Load params
    bytes, _ = ioutil.ReadFile(modParams)
    funp, _ = graphmod.GetFunction("load_params")
    funp.Invoke(bytes)


    // Set module input
    funp, _ = graphmod.GetFunction("set_input")
    funp.Invoke("input", inX)

    // Run or Execute the graph
    funp, _ = graphmod.GetFunction("run")
    funp.Invoke()

    // Get output from runtime.
    funp, _ = graphmod.GetFunction("get_output")
    funp.Invoke(int64(0), out)

    // Access output tensor data.
    outIntf, _ := out.AsSlice()
    outSlice := outIntf.([]float32)

    // outSlice here holds flattened output data as a golang slice.
}

gotvm 扩展了 TVM 打包函数系统，以支持将 golang 函数闭包作为打包函数。示例 可用于将 golang 闭包注册为 TVM 打包函数，并在编程语言边界之间调用它。

显示代码

参考

  • [1] Go 编程语言
  • [2] Go 文档指南
  • [3] Go 测试用例框架
  • [4] Go CFFI
  • [5] Go 可变参数函数
  • [6] CFFI 参考
  • [7] Go Finalizers