【AI芯片】ONNXRuntime
ONNXRuntime是微软推出的一款推理框架,用户可以非常便利的用其运行一个onnx模型。ONNXRuntime支持多种运行后端包括CPU,GPU,TensorRT,DML等。可以说ONNXRuntime是对ONNX模型最原生的支持。
虽然大家用ONNX时更多的是作为一个中间表示,从pytorch转到onnx后直接喂到TensorRT或MNN等各种后端框架,但这并不能否认ONNXRuntime是一款非常优秀的推理框架。而且由于其自身只包含推理功能(最新的ONNXRuntime甚至已经可以训练),通过阅读其源码可以解深度学习框架的一些核心功能原理(op注册,内存管理,运行逻辑等)
import numpy as np
import onnx
import onnxruntime as ort
image = cv2.imread("image.jpg")
image = np.expand_dims(image, axis=0)
onnx_model = onnx.load_model("resnet18.onnx")
sess = ort.InferenceSession(onnx_model.SerializeToString())
sess.set_providers(['CPUExecutionProvider'])
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
output = sess.run([output_name], {input_name : image_data})
prob = np.squeeze(output[0])
print("predicting label:", np.argmax(prob))
总体来看,整个ONNXRuntime的运行可以分为三个阶段,Session构造,模型加载与初始化和运行。和其他所有主流框架相同,ONNXRuntime最常用的语言是python,而实际负责执行框架运行的则是C++。
第一阶段 Session构造
构造阶段即创建一个InferenceSession对象。在python前端构建Session对象时,python端会通过onnxruntime_pybind_state.cc调用C++中的InferenceSession类构造函数,得到一个InferenceSession对象。
InferenceSession构造阶段会进行各个成员的初始化,成员包括负责OpKernel管理的KernelRegistryManager对象,持有Session配置信息的SessionOptions对象,负责图分割的GraphTransformerManager,负责log管理的LoggingManager等。当然,这个时候InferenceSession就是一个空壳子,只完成了对成员对象的初始构建。
第二阶段 模型加载与初始化
在完成InferenceSession对象的构造后,会将onnx模型加载到InferenceSession中并进行进一步的初始化。
1 . 模型加载
模型加载时,会在C++后端会调用对应的Load()函数,InferenceSession一共提供了8种Load函数。包读从url,ModelProto,void* model data,model istream等读取ModelProto。InferenceSession会对ModelProto进行解析然后持有其对应的Model成员。
2 . Providers注册
在Load函数结束后,InferenceSession会调用两个函数:RegisterExecutionProviders()和sess->Initialize();
RegisterExecutionProviders函数会完成ExecutionProvider的注册工作。这里解释一下ExecutionProvider,ONNXRuntime用Provider表示不同的运行设备比如CUDAProvider等。目前ONNXRuntimev1.0支持了包括CPU,CUDA,TensorRT,MKL等七种Providers。通过调用sess->RegisterExecutionProvider()函数,InferenceSession通过一个list持有当前运行环境中支持的ExecutionProviders。
3 . InferenceSession初始化
即sess->Initialize(),这时InferenceSession会根据自身持有的model和execution providers进行进一步的初始化(在第一阶段Session构造时仅仅持有了空壳子成员变量)。该步骤是InferenceSession初始化的核心,一系列核心操作如内存分配,model partition,kernel注册等都会在这个阶段完成。
首先,session会根据level注册 graph optimization transformers,并通过GraphTransformerManager成员进行持有。
接下来session会进行OpKernel注册,OpKernel即定义的各个node对应在不同运行设备上的计算逻辑。这个过程会将持有的各个ExecutionProvider上定义的所有node对应的Kernel注册到session中,session通过KernelRegistryManager成员进行持有和管理。
然后session会对Graph进行图变换,包括插入copy节点,cast节点等。
接下来是model partition,也就是根运行设备对graph进行切分,决定每个node运行在哪个provider上。
最后,为每个node创建ExecutePlan,运行计划主要包含了各个op的执行顺序,内存申请管理,内存复用管理等操作。
第三阶段 模型运行
模型运行即InferenceSession每次读入一个batch的数据并进行计算得到模型的最终输出。然而其实绝大多数的工作早已经在InferenceSession初始化阶段完成。run阶段主要是顺序调用各个node的对应OpKernel进行计算。
更多推荐
所有评论(0)