为 OneFlow 添加新的前端语言

因为各种机缘巧合和历史的必然， Python 成为了现在事实意义上的“人工智能编程语言”， OneFlow 也把 Python 作为了用户接口语言。然而事实上，Python 只是 OneFlow 的前端，复杂的计算和并行功能代码，还是通过 C/C++ 实现的，OneFlow 良好的解耦设计，使得我们可以较容易的支持 Python 作为用户接口，自然也可以支持更多的语言作为用户接口。

在这个项目中，我们将给 OneFlow 添加 Java 前端，支持模型加载和模型推理的功能。有了模型加载和推理功能，用户可以很容易在自己的 Java 应用中加载训练好的模型，将模型部署上线！

在写项目申报书，写着写着突然意识到了这个问题。深度学习框架支持 Java 前端有什么意义呢？模型分开部署不香吗？后来在和一些业内人士进行了交流之后，才想明白。首先，有一些老业务用的还是 Java，如果要继承过来，就必须要考虑 Java。其次，将模型分开部署，需要搭建服务，用 HTTP 或者 RPC 的方式来调用服务往往有网络上的延迟，然而给某个服务的时间总共才 10 ms，所以要考虑如何减少网络通信带来的延迟。为了构建一个满足响应时间的服务，就需要考虑如何将模型内嵌到现有的应用当中。

计划阶段，需要考虑采用什么技术，梳理整个程序流程。

OneFlow 底层是通过 C++ 来实现的，因此可以调用动态链接库来为 OneFlow 添加一门新的前端语言。对于 Java，可以通过 Java Native Interface 来调用本地方法，所以在我们的这个项目中，使用 Java Native Interface 作为一个调用的桥梁，沟通 C++ 和 Java。

确定了使用 JNI 之后，接下来需要做的是理清楚调用流程。OneFlow 是如何启动的？如何加载模型？如何前向传播？如何输入，又如何输出呢？带着这些问题，去阅读源代码吧。

我们需要做的是加载模型和进行前向传播，获取推理的结果。于是，我将主要精力放在了阅读 InferenceSession.py 这个代码文件上。此外 OneFlow 的研发工程师 @Lyon 在知乎上分享了三篇关于 “一个 Job 在 OneFlow 中的执行过程” ，这三篇文章对整个流程有很详细的分析。不断阅读代码，单步调试，不断地单步跟进去，从 Python 调试到 C++ 底层，最后总算是对整个流程有了清晰的认识。

于是，我将整个流程分为 5 个阶段，分别是：初始化阶段，模型加载和编译阶段，启动阶段，推理阶段，关闭阶段。

初始化阶段

在 Python 中，import oneflow 背后做了一些初始化的工作，这部分的工作切不可忽略。这部分完成的工作有：初始化物理环境，初始化默认的 Session，设置运行模式，注册结束时调用的函数。接着就是在 InferenceSession 中可以看到的初始化：env 初始化，scope 初始化，session 初始化。值得注意的是：scope 初始化，在 Python 中出现了函数闭包作为参数传给 C++，并且使用 nonlocal 来获取运行的结果。这个操作，真是太骚了。

模型加载和编译阶段

读取保存在本地的 protobuf 文件，设置检查点 (checkpoint)，编译计算图。

调用 C++ 的接口，StartLazyGlobalSession，读取检查点，加载模型权重。加载权重是这个部分的难点，需要定义一个 JobInstance，传递函数对象给 JobInstance，然后启动一个 Job。

输入使用 Push Job 来推送数据，然后启动一个 Job 来进行前向传播，最后输出使用 Pull Job 来拉取数据。

这部分的难点在于 Tensor 类的设计。这个部分对比参考了同为深度学习框架的 Pytorch 和 MindSpore，这两者都有 Java 前端。最后的设计是，在 Java 端保存数据，而不是保存一个指针。这样的好处就是，用户无需关注内存的状况，用户不需要显式调用一个方法来清理 Tensor 中指针指向的内容，让 GC 去担心就好了。

调用 stopLazyGlobalSession 和 destroyLazyGlobalSession 即可。

整个实施过程，大致可以划分为几个阶段。

最初，我不会 CMake 的时候，不知道如何编译一个动态链接库。于是，我有了这样的想法：在 Java 生成的 native 实现中，使用 C 去调用 pybind11 生成的动态链接库。链接的时候，还要把 libpython.so 带上。emmm... 似乎笨了点，但这种探索本身却是很有意思的，哈哈。

后来学了一点 CMake，直接生成 Java 需要的动态链接库就好了。于是 Java 和 C++ 交互的问题就搞定了。

面向功能编程

设计什么的都先甩开！面向功能编程，gkd。

初始化，加载编译，启动，推理，关闭。一套整搞下来之后，将 Java 推理的结果和 Python 推理的结果一比，不一样啊！经过一段苦苦的 Debug 之后，发现原来是大端小端的问题！Java 默认是大端编码的，而 x86 系列的 CPU 都是小端编码的。

修改编码问题之后，结果一比。啊！终于，终于，终于，完成功能了。

第一次整理代码前，导师还跟我开了个腾讯会议。看着这乱七八糟的代码，和老师汇报了一下进度。有些地方甚至忘记了为什么那么写，实在是惭愧。我真是太菜了。这次开会老师和我特别提到，要尽可能减少内存的复制。

性能优化阶段

当时面向功能编程的第一版的实现中，使用 Get ArrayElements Routines 或者 Get ArrayRegion Routines 来获取数组。根据 JVM 具体的实现，可能发生数组的复制。

对于这个问题，可以使用 DirectBuffer。它是 Java 堆外内存，是一块连续的内存，在 DirectBuffer 整个对象被 GC 之前，这个地址和对应的内容是稳定不变的， 这个 StackOverflow 的问答有一定的启发性 。使用 DirectBuffer 的好处是，避免数组的复制。根据 JNI 的文档， GetDirectBufferAddress 这个方法可以直接访问到那个地址。

于是我重新设计了一下 Tensor 类，参考对比了 Pytorch 和 MindSpore 的设计。最终决定：将数据保存在 Java 这边，并且保存在 DirectBuffer 中这个方案。毕竟，咱用 Java 的人不想关注内存呢，专门写个方法来释放 C++ 的数据，不够优雅。

思考设计，重构

面向功能编程，好处是最快的速度完成功能，以获得成就感。可是，没有设计会导致扩展的困难。在优化的过程中，也发现了修改代码很麻烦。于是我认为，是时候思考设计了。我的设计很简单，将代码分层，使得容易扩展。native 代码，分为两个部分，一部分是 OneFlow 实际执行的逻辑，一部分是和 Java 交互的逻辑。这两种逻辑分开之后，扩展就方便了，而且 OneFlow 实际执行的逻辑还可以作为 C API 复用。Java 部分的代码，需要考虑用户使用的友好程度，如何设计一个用户友好的接口很关键。

添加新功能

signature 和 batching。前期的方案中忽略了这两个东西，在重新设计之后，发现新增加这两个功能，异常的简单，只需要增加几行代码就可以完成了。

不过在这几行代码之前，有个小插曲。signature 实现过程中遇到了一个小问题，最终定位到了编译图时候的一个 Pass，有个地方忘了检查 map 键是否存在。于是提了个 PR，修复了这个问题。然而我真的太菜了，没有去考虑一下代码的上下文，就检查 key，然后直接跳过。实际上，往后面看几行，会发现，那么写会导致查找两次。

编写 CMake 代码

之前一直在 OneFlow 的仓库之外构建 Jar 包，后面老师说要放到 OneFlow 内部。为了编译 jar 包，需要解决 jar 包的依赖。这个 jar 依赖了 proto 对应的 java 文件，以及 protobuf-java.jar。继续学一点 CMake，然后构建依赖，最后构建 Jar 包。完善 CMake 代码前，整个项目的构建很不方便，需要自己手动编译 proto 为 java 文件，还要手动修改一个 message 的名字以避免 protobuf 的 BUG (后面发现是版本问题)。完善 CMake 代码后，构建就很方便了，只需要几行 make 命令就可以完成构建。

再一次重构

为了减少 Java 和 C++ 之间序列化的开销，决定 Java 端不再使用 protobuf。这一次重构，大部分 Java 代码改为了 C++ 代码，于是 C++ 接口的 API 粒度更大一点，Java 只需要简单调用一下，不需要复杂的逻辑了。

于是，我发现了，前面几天辛辛苦苦写的 CMake 代码作废了。。。清理一下代码吧 😦

接下来还要不断完善代码，每每阅读 OneFlow 内部优雅的代码，就越发觉自己的代码很矬，再努力努力吧。

这段开源之夏的经历，更像是一种探索，打代码的过程非常的开心和快乐！

这次活动促进了开源软件的发展和优秀开源软件社区建设，增加开源项目的活跃度，推进开源生态的发展；感谢开源之夏主办方为这次活动提供的平台与机会。感谢导师在这过程中对我的付出和指导；感谢 OneFlow 社区，他们提出的 Actor + SBP 的设计，让我大受震撼，感谢~