简介

ONNX Runtime 是一个用于 ONNX(Open Neural Network Exchange) 模型推理的引擎。微软联合 Facebook 等在 2017 年搞了个深度学习以及机器学习模型的格式标准 —ONNX，顺路提供了一个专门用于 ONNX 模型推理的引擎，onnxruntime。目前 ONNX Runtime 还只能跑在 HOST 端，不过官网也表示，对于移动端的适配工作也在进行中。
一半处于工作需要一半出于兴趣，决定阅读一下 onnxruntime 的源码。这里算个学习记录吧。

安装

ONNX Runtime 的 GitHub 仓库地址为 https://github.com/microsoft/onnxruntime 。编译安装过程可以参照 GitHub 上的说明，这里为了方便，直接选择了 PyPi 的安装源。执行

即完成了安装。需要注意的是只支持 Python3。

开始

涉及文件

onnxruntime\onnxruntime\python\session.py
onnxruntime\onnxruntime\core\framework\utils.cc
onnxruntime\onnxruntime\python\onnxruntime_pybind_state.cc
onnxruntime\onnxruntime\core\session\inference_session.cc
onnxruntime\onnxruntime\core\session\inference_session.h

代码入口

代码阅读需要先找到一个入口。通过 onnxruntime 的例子我们知道，在 Python 使用使用 onnxruntime 很简单，主要代码就三行：

  1. import onnxruntime
  2. sess = onnxruntime.InferenceSession('YouModelPath.onnx')
  3. output = sess.run([output_nodes], {input_nodes: x})

第一行导入 onnxruntime 模块；第二行创建一个InferenceSession的实例并传给它一个模型地址；第三行调用run方法进行模型推理。因此 onnxruntime 模块中的InferenceSession就是我们的切入点。

实例生成

ONNX Runtime 的代码组织非常良好，我们很容易找到InferenceSession所在文件session.py，整个文件非常简单，就只定义了一个InferenceSession类。通过阅读InferenceSession的__init__函数，

  1.  def __init__(self, path_or_bytes, sess_options=None, providers=[]):
  2.         """
  3.         :param path_or_bytes: filename or serialized model in a byte string
  4.         :param sess_options: session options
  5.         :param providers: providers to use for session. If empty, will use
  6.             all available providers.
  7.         """
  8.         self._path_or_bytes = path_or_bytes
  9.         self._sess_options = sess_options
  10.         self._load_model(providers)
  11.         self._enable_fallback = True
  13.     def _load_model(self, providers=[]):
  14.         if isinstance(self._path_or_bytes, str): 
  15.             self._sess = C.InferenceSession(
  16.                 self._sess_options if self._sess_options else C.get_default_session_options(), 
  17.                 self._path_or_bytes, True)
  18.         elif isinstance(self._path_or_bytes, bytes):
  19.             self._sess = C.InferenceSession(
  20.                 self._sess_options if self._sess_options else C.get_default_session_options(), 
  21.                 self._path_or_bytes, False)
  25.         else:
  26.             raise TypeError("Unable to load from type '{0}'".format(type(self._path_or_bytes)))
  28.         self._sess.load_model(providers)

我们发现其实这里InferenceSession只不过是一个壳，所有工作都委托给了C.InferenceSession，C从导入语句from onnxruntime.capi import _pybind_state as C可知其实就是一个 C 语言实现的 Python 接口，其源码在onnxruntime\onnxruntime\python\onnxruntime_pybind_state.cc中。onnxruntime_pybind_state.cc是将 C++ 代码暴露给 Python 的一个接口，就像是一个门，代码经过这里，就从 Python 进入了 C++ 的世界。
门在这了，开门 的钥匙在哪儿？
我们进盯着 Python 中

  1.  self._sess = C.InferenceSession(
  2.                 self._sess_options if self._sess_options else C.get_default_session_options(), 
  3.                 self._path_or_bytes, True)

这句话，它是全村的希望。通过这句话，我们知道，在onnxruntime_pybind_state.cc应该会定义有一个类，名叫InferenceSession，一顿操作猛如虎，定位到InferenceSession定义的地方：

  1. py::class_<InferenceSession>(m, "InferenceSession", R"pbdoc(This is the main class used to run a model.)pbdoc")
  5.       .def(py::init([](const SessionOptions& so, const std::string& arg, bool is_arg_file_name) {
  7.         if (is_arg_file_name) {
  8.           return onnxruntime::make_unique<InferenceSession>(so, arg, SessionObjectInitializer::Get());
  9.         }
  12.         std::istringstream buffer(arg);
  13.         return onnxruntime::make_unique<InferenceSession>(so, buffer, SessionObjectInitializer::Get());
  14.       }))

欢迎来到 C++。def(py::init([](const SessionOptions& so, const std::string& arg, bool is_arg_file_name)实现了类似 Python 中__init__的功能，其根据传入的模型参数类型（模型的地址还是模型的数据流），调用 C++ 中的类InferenceSession的相应构造函数构造一个的实例，然后将这个实例的指针返回给 Python。由于我们例子中传入的是模型的地址字符串，因此我们需要找到的是签名类型为：

  1. InferenceSession(const SessionOptions& session_options,
  2.                    const std::string& model_uri,
  3.                    logging::LoggingManager* logging_manager = nullptr);

的构造函数。
这里有个奇怪的现象：

  1.  if (is_arg_file_name) {
  2.           return onnxruntime::make_unique<InferenceSession>(so, arg, SessionObjectInitializer::Get());
  3.         }

中第三个参数我们通过查看SessionObjectInitializer::Get()获取到的是类SessionObjectInitializer的一个实例，但是InferenceSession对应的构造函数对应为所需要的是一个logging::LoggingManager的指针，对不上，咋整？我们知道 C++ 可不像 Python，C++ 是强类型的语言，不将就。这里作者用了个小技巧，他为SessionObjectInitializer定义了两个类型转换函数，让编译器帮他转到所需要的类型，这里编译器会将SessionObjectInitializer转换成logging::LoggingManager*。
来看看

  1. InferenceSession(const SessionOptions& session_options,
  2.                    const std::string& model_uri,
  3.                    logging::LoggingManager* logging_manager = nullptr);

的实现：

  1. InferenceSession::InferenceSession(const SessionOptions& session_options,
  2.                                    const std::string& model_uri,
  3.                                    logging::LoggingManager* logging_manager)
  4.     : insert_cast_transformer_("CastFloat16Transformer") {
  5.   model_location_ = ToWideString(model_uri);
  6.   model_proto_ = onnxruntime::make_unique<ONNX_NAMESPACE::ModelProto>();
  7.   auto status = Model::Load(model_location_, *model_proto_);
  8.   ORT_ENFORCE(status.IsOK(), "Given model could not be parsed while creating inference session. Error message: ",
  9.               status.ErrorMessage());
  12.   ConstructorCommon(session_options, logging_manager);
  13. }

这里主要就做了三件事：

  1. 将模型地址保存在类成员变量model_location_中；
  2. 将模型二进制内容保存在类成员变量model_proto_;
  3. 调用ConstructorCommon完成剩余的工作。
    ConstructorCommon中做些环境检查，准备 log 输出等工作。其中最主要的是，是创建了一个SessionState实例session_state_，这是类成员变量，其中打包了为运行这个模型所需要的线程池、模型结构、provider 等信息。至于什么是 Provider，其实就是模型所跑的硬件，比如是 CPU 还是 GPU，到了这里其实session_state_里面很多信息还没完备，例如模型结构并未保存，Provider 还只是个壳，里面并没有保存任何硬件信息，还需要一个初始化阶段。至此，InferenceSession实例创建完毕。

初始化

又回到最初的起点，Python 代码开始的地方，最后一句self._sess.load_model(providers)，其实现如下：

  1. .def(
  2.           "load_model", [](InferenceSession* sess, std::vector<std::string>& provider_types) {
  3.             OrtPybindThrowIfError(sess->Load());
  4.             InitializeSession(sess, provider_types);
  5.           },
  6.           R"pbdoc(Load a model saved in ONNX format.)pbdoc")

load_model主要做了一下事情：

  1. 将模型二进制内容解析；
  2. 选择模型运行方式，并行还是串行；
  3. 选择模型 Provider，如果用户没有指定 Provider，就把目前运行环境中支持的硬件都注册，比如 GPU，CPU 等，并且保证 CPU 一定可用；
  4. 确定模型中各个节点的运行先后顺序。
    这里先不细说了，只需要知道它是按照 ONNX 标准将二进制数据解析成一个图并将它存储在session_stat_中就可以了。以后再详细说。经过这一步之后，session_state_已经完备，到达神装，可以随时开战。

运行

经过初始化之后，一切就绪。我们直接看 C++ 中InferenceSession的run方法好了，因为通过前面知道，在 Python 中的操作最终都会调用到 C++ 的代码来执行实际的内容。虽然InferenceSession重载了很多run方法，但是最终都会辗转调用到签名为

  1. Status InferenceSession::Run(const RunOptions& run_options, const std::vector<std::string>& feed_names,
  2.                              const std::vector<OrtValue>& feeds, const std::vector<std::string>& output_names,
  3.                              std::vector<OrtValue>* p_fetches)

的这个。在这里，run方法对输入数据做了些检查等工作后，变将数据、模型信息，provider 信息等，传递给了utils::ExecuteGraph:

  1. utils::ExecuteGraph(*session_state_, feeds_fetches_manager, feeds, *p_fetches,
  2.                             session_options_.execution_mode,
  3.                             run_options.terminate, run_logger))

而utils::ExecuteGraph反手又将工作委托给了utils::ExecuteGraphImpl，而utils::ExecuteGraphImpl将会根据前面初始化中确定的各个 node 的执行先后顺序，找到 node 类似对对应的 kernel，调用他们Compute()方法进行计算。

总结

一个大概流程就是通过使用 pybind11 将 C++ 接口暴露给 Python，Python 经过简单封装后提供给用户直接使用。上面有几个关键点值得深入研究：

  1. 模型节点执行顺序的确定；
  2. 模型节点 Provider 的选择；
  3. 模型解析过程；
  4. 模型推理详细过程；
  5. 模型如何高效推理。
    最后，一图胜千言：
    ONNX Runtime 源码阅读：模型推理过程概览 - 简书 - 图1
    onnxruntime_exec_flow.png
    https://www.jianshu.com/p/c50953e8d828