内存优化¶
模型 Inference 时候需要占用处理器上面的内存用于代码执行,数据缓存等,目前处理器上的内存(RAM)通常都比较大,但是在一些情况下内存依然是瓶颈,如:
在 NVIDIA 设备上进行多路视频解码和检测时候,会被 CUDA device 上的显存大小限制同时开启的路数。
在 TEE 环境中,内存资源非常受限。
以及在一些特定的嵌入式平台上,为了成本考虑一般内存有限。
MegEngine Lite 对 Inference 阶段的内存进行分类,主要包含下面3类:
存放模型权重的内存 :用于加载模型中的权重(weights)数据,并在 Inference 阶段参与计算。
计算图结构内存 :模型 Load 之后保存其结构信息所需的内存。
Runtime 内存 :模型计算过程中用于保存 Operator 输入输出的 Tensor 内存,以及具体 Operator 计算时候需要的临时内存。
MegEngine Lite 中默认已经对静态图的内存做了极致的内存规划,如果在这个极致的内存规划上还需要继续进行内存优化,MegEngine lite 还有下面的方法进行内存优化:
Record2 内存优化
共享运行时内存
多 Network 共享权重
用户注册自定义内存池
Record2 内存优化¶
Record2 是在 Record 优化: 基础上进一步优化内存,原理为:在 record1 的基础上, 会析构掉原来模型 load 时候创建的计算图,即上面提到的 计算图结构内存 。它和 record1 有相同的 限制。
Warning
record 主要有如下3个 限制条件 :
执行的模型必须所有 shape 是静态可推导的,不能有动态图情况。
输入 Tensor 的 shape 不能改变,改变 shape 之后,所有的内存计划都会改变,导致记录的 kerns 不可用。
模型中只能在一个设备上进行,目前 Record 功能只支持 CPU 上模型执行。
配置方法:将 基本的配置代码 注释 1 的地方,将 Option 中的 comp_node_seq_record_level,设置为 2 即 record2,参考:创建 Network, 另外开启 record2 还需要设置 Option 中 var_sanity_check_first_run 为 false, fake_next_exec 为 true。
Note
record2 内存节省的大小需要在直接模型和环境中测试。
共享运行时内存¶
多个 NetWork 不同时运行时候,可以共享他们的 Runtime 内存,这样可以有效减少多个模型的 Inference 的场景下内存的峰值。大概原理为,MegEngine 会为多个模型分配同一个 Runtime 内存,这个内存不是所有共享模型的 Runtime 内存的和,而是取他们的最大值。
配置方法: 调用 MegEngine Lite 的 Runtime API 接口,参考 share_runtime_memory_with。
Config config;
std::shared_ptr<Network> network_src = std::make_shared<Network>(config);
std::shared_ptr<Network> network_dst = std::make_shared<Network>(config);
network_src->load_model(model_path_src);
Runtime::share_runtime_memory_with(network_dst, network_src);
network_dst->load_model(model_path_dst);
network_src = LiteNetwork()
network_src.load(model_path_src)
network_dst = LiteNetwork()
network_dst.share_runtime_memroy(src_network)
network_dst.load(model_path_dst)
上面是在 CPP 和 python 中使用共享运行时内存优化的代码片段,主要步骤为:
创建源 Network 之后,可以进行正常的加载模型等操作。
创建目的 Network。
调用 share_runtime_memroy 接口配置目的 Network 的 Runtime 内存将从源 Network 中共享。
目的 Network 进行模型加载等操作。
Warning
目的 Network 设置共享 Runtime 内存之前不能加载模型,源 Network 没有任何限制。
目的 Network 和源 Network 不能同时运行,否则出现计算错误,这需要 用户自己保证。
多 Network 共享权重¶
当需要对同一个模型创建多个 Network 时候,如果每个 Network 都载入一次模型,这样势必会将模型的权重拷贝多次,导致存放模型权重的内存成倍的增大。 MegEngine Lite 支持这种情况下多个 Network 共享同一份模型的权重,这样将有效的减少内存的用量,典型的应用场景是,同一个模型希望在多个线程中 同时并行进行 Inference,这时就可以让这些不同线程里面的 Network 共享同一份权重数据,但是他们 Network 是不同的。
配置方法: 调用 MegEngine Lite 的 Runtime API 接口参考 shared_weight_with_network。
std::shared_ptr<Network> network_src = std::make_shared<Network>();
network->load_model(model_path);
std::shared_ptr<Network> network_dst = std::make_shared<Network>(config);
Runtime::shared_weight_with_network(network_dst, network_src);
network_src->forward();
network_src->wait();
network_dst->forward();
network_dst->wait();
network_src = LiteNetwork()
network_src.load(model_path_src)
network_dst = LiteNetwork()
network_dst.share_weights_with(src_network)
network_src.forward()
network_src.wait()
network_dst.forward()
network_dst.wait()
上面是在 CPP 和 python 中多个模型共享权重的代码片段,主要步骤为:
创建源 Network 之后,创建目的 Network。
源 Network 进行模型加载,推理等操作
调用 share_weights_with 接口将目的 Network 的权重内存从源 Network 中共享过去。
目的 Network 进行模型加载等操作。
Note
目的 Network 不需要再载入模型
目的 Network 和源 Network 可以同时运行
用户注册自定义内存池¶
用户可以在 MegEngine Lite 中为模型运行时的内存分配注册回调函数,这样 Runtime 期间 MegEngine Lite 申请的内存都可以被外部回调函数接管。 某些场景下,用户可以在外部实现一个内存池,可以使得 MegEngine Lite 申请的内存和外部代码申请的内存进行复用,减少用量内存的峰值,目前该接口只能 在 C++ 接口中使用。
MegEngine Lite 中定义了 Allocator 的基类,用户可以继承这个基类,然后 override 它的 allocate 和 free 接口。
allocate 接口:给定了需要分配内存所在的设备类型,设备的 ID,以及需要分配内存的长度,内存对齐的要求,需要以 void* 的形式返回申请好的内存。
free 接口:传递进来需要释放内存的设备类型,设备 ID,以及需要释放的内存指针。
class Allocator {
public:
virtual ~Allocator() = default;
//! allocate memory of size in the given device with the given align
virtual void* allocate(
LiteDeviceType device_type, int device_id, size_t size, size_t align) = 0;
//! free the memory pointed by ptr in the given device
virtual void free(LiteDeviceType device_type, int device_id, void* ptr) = 0;
};
下面是使用 MegEngine Lite 进行推理,并使用用户自定义的内存分配器进行 Runtime 的内存分配的 example
class MyAllocator : public lite::Allocator {
public:
//! allocate memory of size in the given device with the given align
void* allocate(LiteDeviceType, int, size_t size, size_t align) override {
return memalign(align, size);
};
//! free the memory pointed by ptr in the given device
void free(LiteDeviceType, int, void* ptr) override {
free(ptr);
};
};
auto allocator = std::make_shared<MyAllocator>();
//! create and load the network
std::shared_ptr<Network> network = std::make_shared<Network>();
Runtime::set_memory_allocator(network, allocator);
network->load_model(network_path);
//! forward
network->forward();
network->wait();
为了完成自定义内存申请,用户需要:
自定义一个 MyAllocator 并继承自 lite::Allocator。
创建一个 std::shared_ptr<MyAllocator> 的智能指针。
创建完成 Network 之后,调用 Runtime 的 set_memory_allocator 将自定义的 MyAllocator 注册 MegEngine Lite。
正常运行模型。
Warning
MegEngine Lite 中只有 Runtime 阶段使用的内存才可能使用到用户自定义的分配器。