mxnet 源码阅读 5

mxnet/src/storage/storage_manager.h

mxnet/src/storage/pooled_storage_manager.h

 

storage_manager.h 中的类StorageManager是抽象的接口层，以纯虚函数 Alloc Free DirectFree 给出，该类属于mxnet 下的storage名字空间。

GPUPooledStorageManager 继承自StorageManager，在其构造函数中通过获取环境变量MXNET_GPU_MEM_POOL_RESERVE设置GPU内存保留大小百分比，如果该参数没有设置则默认为5

同时析构函数释放所有GPU 显存。

重载的Alloc Free 和DirectFree 通过unordered_map< size, vector<void*> > 以内存大小作为键，内存地址vector作为值，例如 5M的地址对应 0x00002 0x00007 等。

Alloc函数执行时先检查内存池中是否有该大小的内存，如果有则返回该地址，并且在内存池中移除该地址；如果没有则申请之。

Free函数直接将要free的地址添加到内存池map中。

当然该类必然有接口可以实现内存的释放，DirectFree和析构函数都可以完成该任务，其中DirectFree释放给定地址的内存空间，析构函数调用ReleaseAll释放所有内存池。

 

上面所有的内存都是指GPU显存。