查看原文
其他

为OneFlow添加新的前端语言

OneFlow 2022-05-10

The following article is from 开源之夏 Author 周泽楷


撰文 | 周泽楷


在近期举办的开源之夏“暑期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新一代开源深度学习框架



您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存