本文探讨了TensorFlow中GPU的BFC（Best Fit with Coalescing）allocator的工作原理，包括其创建过程、显存分配策略以及如何查看当前实际使用显存量。BFC allocator在显存不足时会翻倍申请，可能导致实际使用小于申请量的情况。通过修改分配策略或使用特定Op（MaxBytesInUse）可以更有效地监控和管理GPU显存。
主要写GPU 的bfc allocator，业务需要稍微调研了一下，写个大概，很多细节没有去看。
一、源码
1、BaseGPUDeviceFactory::CreateGPUDevice中使用GPUProcessState单例调用process_state->GetGPUAllocator 创建GPU allocator
2、GetGPUAllocator 实例化allocator
其中sub_allocator 也就是GPUMemAllocator是通过stream去调用GPU alloc显存的接口，也就是实际去申请GPU显存的allocator，sub_allocator由GPUBFCAllocator控制使用。GPUBFCAllocator是BFCAllocator的子类，通过多态BFCAllocator的基类指针指向GPUBFCAllocator对象。
3、显存的申请策略
简要说明一下实际申请显存的策略，具体的显存分配算法没有看。
显存的申请函数是AllocateRaw，释放函数是DeallocateRaw。AllocateRaw的策略是直接通过sub_allocaotr申请大块的显存，然后通过分配算法将合适的一块显存首地址返回，如果剩下的显存不够分配了，再通过sub_allocator直接申请翻倍的显存，如当前已有1G的显存，用完后会再申请1G，变成2G，因此显存的实际申请是1G->2G->4G->8G->16G， 一直到最后剩下可申请的全部申请完，如果显卡显存不够了，且TF的一
最低0.47元/天 解锁文章
扫一扫
686
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
