为OneFlow添加新的前端语言

撰文 | 周泽楷

在近期举办的开源之夏“暑期2021”活动中，来自OneFlow社区的开发者周泽楷分享了“为OneFlow添加新的前端语言”的项目经验。

1

简介

任务介绍

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

项目目标

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

项目意义

在写项目申报书，写着写着突然意识到了这个问题。深度学习框架支持 Java 前端有什么意义呢？模型分开部署不香吗？后来在和一些业内人士进行了交流之后，才想明白。

首先，有一些老业务用的还是 Java，如果要继承过来，就必须要考虑 Java。其次，将模型分开部署，需要搭建服务，用 HTTP 或者 RPC 的方式来调用服务往往有网络上的延迟，然而给某个服务的时间总共才 10ms，所以要考虑如何减少网络通信带来的延迟。为了构建一个满足响应时间的服务，就需要考虑如何将模型内嵌到现有的应用当中。

2

计划

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

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

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

我们需要做的是加载模型和进行前向传播，获取推理的结果。于是，我将主要精力放在了阅读InferenceSession.py

https://github.com/Oneflow-Inc/oneflow/blob/master/python/oneflow/serving/inference_session.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（参考链接：

https://github.com/pytorch/pytorch/blob/master/android/pytorch_android/src/main/java/org/pytorch/Tensor.java） 

和 MindSpore（参考链接：

https://gitee.com/mindspore/mindspore/blob/r1.2/mindspore/lite/java/java/common/src/main/java/com/mindspore/lite/MSTensor.java），

这两者都有 Java 前端。最后的设计是，在 Java 端保存数据，而不是保存一个指针。这样的好处就是，用户无需关注内存的状况，用户不需要显式调用一个方法来清理 Tensor 中指针指向的内容，让 GC 去担心就好了。

关闭阶段

调用 stopLazyGlobalSession 和 destroyLazyGlobalSession 即可。

3

实施

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

• 探索阶段，主要解决 Java 和 C++ 交互的问题。

• 面向功能编程：初始化，加载编译，启动，推理，关闭。

• 整理代码。

• 性能优化阶段：设计 Tensor 类，尽可能避免内存复制。

• 思考设计，重构。将代码分层，使得容易扩展。

• 添加新功能 signature 和 batching。

• 编写 CMake 代码。构建 jar 包和动态链接库。

• 再一次重构。Java 端不再使用 protobuf。

探索阶段

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

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

面向功能编程

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

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

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

整理代码

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

性能优化阶段

当时面向功能编程的第一版的实现中，使用 GetArrayElementsRoutines

(https://docs.oracle.com/javase/8/docs/technotes/guides/jni/spec/functions.html#Get_PrimitiveType_ArrayElements_routines)

或者 GetArrayRegion Routines

(https://docs.oracle.com/javase/8/docs/technotes/guides/jni/spec/functions.html#Get_PrimitiveType_ArrayRegion_routines) 来获取数组。根据 JVM 具体的实现，可能发生数组的复制。

对于这个问题，可以使用 DirectBuffer。它是 Java 堆外内存，是一块连续的内存，在 DirectBuffer 整个对象被 GC 之前，这个地址和对应的内容是稳定不变的，这个 StackOverflow 的问答有一定的启发性

(https://stackoverflow.com/questions/1854398/how-to-garbage-collect-a-direct-buffer-in-java)。

使用 DirectBuffer 的好处是，避免数组的复制。根据 JNI 的文档，GetDirectBufferAddress

(https://docs.oracle.com/javase/8/docs/technotes/guides/jni/spec/functions.html#:~:text=This%20function%20allows%20native%20code%20to%20access%20the%20same%20memory%20region%20that%20is%20accessible%20to%20Java%20code%20via%20the%20buffer%20object.) 

这个方法可以直接访问到那个地址。

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

思考设计，重构

面向功能编程，好处是最快的速度完成功能，以获得成就感。可是，没有设计会导致扩展的困难。在优化的过程中，也发现了修改代码很麻烦。于是我认为，是时候思考设计了。我的设计很简单，将代码分层，使得容易扩展。

native 代码，分为两个部分，一部分是 OneFlow 实际执行的逻辑，一部分是和 Java 交互的逻辑。这两种逻辑分开之后，扩展就方便了，而且 OneFlow 实际执行的逻辑还可以作为 C API 复用。Java 部分的代码，需要考虑用户使用的友好程度，如何设计一个用户友好的接口很关键。

添加新功能

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

不过在这几行代码之前，有个小插曲。signature 实现过程中遇到了一个小问题，最终定位到了编译图时候的一个 Pass，有个地方忘了检查 map 键是否存在。

于是提了个 PR，修复了这个问题。然而我真的太菜了，没有去考虑一下代码的上下文，就检查 key，然后直接跳过。实际上，往后面看几行，会发现，那么写会导致查找两次。

编写 GMake 代码

之前一直在 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 内部优雅的代码，就越发觉自己的代码很矬，再努力努力吧。

4

结语

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

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

题图源自Pexel

其他人都在看

• “我们决定去登月”| OneFlow开源这一年

• 并行机缔造者希利斯和思维机器的浮沉十年

• 对抗软件系统复杂性③：恰当分层，不多不少

• 动态调度的“诅咒”| 原有深度学习框架的缺陷③

• 徐之秋：从游戏启蒙的00后AI工程师

点击“阅读原文”，欢迎下载体验OneFlow新一代开源深度学习框架