RL 使用Cpp动态库加速环境模拟

背景

最近工作在做RL相关的事情, 遇到了环境模拟占用大量时间的问题

性能瓶颈

从业务抽象出RL环境模拟, 如果用Golang计算10000条数据只需要44秒,但由于机器学习需要用pytorch, 直接用python模拟计算10000条数据则需要2000多秒

这种计算性能用在训练时间是难以接受的, 因此需要针对环境计算场景进行优化.

性能分析

经过时间分布占比发现, 在使用python多进程开发后, 有一半多的时间都在环境模拟上, 并且数据集增长会让该模拟时间更快增长. 而这一部份是纯CPU计算, 可以尝试用更快的方法进行计算.

由于业务特殊性, 环境模拟时要进行未来模拟, 因此实际环境模拟的时间可能还要再乘一个数量级的时间 = env_one_time * 10 * e (e为常数)

对于另一部份网络训练的优化可能就需要使用分布式训练, 这一部份不在本章节讨论范围

RoadMap

1. 确认方案

以上旧的方案已经使用多进程尝试对python进行优化, 但仍然不及预期, 因此考虑使用C++动态库的形式给python加速计算

我们需要准备

• python

• g++ or clang

• cmake

2. 保证结果准确性

• 首先需要对c++进行编程, 保证环境模拟时, 与python结果一致

• 产出一份testling, 用Go和python的结果Benchmark, 后续新需求新改动以此为对照

• 使用ipynb做到结果自动对比

• 编写c++单元测试

3. 动态库对接

• 使用cpython对接, python声明相关对应结构题和方法

• 在python加载动态库和方法

• 编写python单元测试

4. 数据对接

• 为了自动内存回收, 避免一些内存泄漏问题, 我们可以在python声明C++需要读和写的数据

• 用C++操作和更新值, C++内部临时变量和指针只能自己申请和释放

• python对接cpp动态库, 一次性初始化数据, 持久化保存创建的数据 (重要)

5. 编译与运行

• c++中新建build目录 cmake CMakefile.txt && make

• python运行RL相关代码 python evalution.py --cfg config/local.json

• 记录时间戳, 测试运行时间

Benchmark

code_version	idx_cnt (10 core)	used time (s)
go	1000	5
python0303	1000	89
py+cpp	1000	15
go	10000	44
python0303	10000	2000+
py+cpp	10000	118

遇到的问题

内存回收

python有时需要传输多维数组给c++进行计算, 而最快速的方法是声明指针,让c++直接按位操作, 但是python的列表数据在底层不像c++那样顺序分布，头指针通过+1无法拿到下一个数据

因此需要用c_types和numpy去转化数据，并用python存储指针

• 而在实际操作后，发现指针传入后，cpp读取不到实际的数据，并且在+n后的地址，读取的数据是正负值都有的随机值，可以猜测到是数据空间被重新初始化了，或者被内存回收了

所以猜测到python在函数内做了数据赋值后，进行了垃圾回收，只有指针引用是不够的，需要将底层数据存储下来

因此最终解决方案是在声明数据后，用memory_keeper将这部分数据进行全局持久化保存，就可以解决这个问题

源码解析

项目已经经过脱敏, 抽象出关键的加载动态库的部分开源到了Github, 有需要的小伙伴可以看看