本文共 1487 字，大约阅读时间需要 4 分钟。
主要写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的一些策略（如通过host 内存来支持更多一点的虚拟显存策略）也无法搞定了，就会出现OOM错误，进程直接结束。
因此我们也可以发现，当nvidia-smi显示占了16G时，实际用的显存可能是因为之前的8G不够了，如果需要8.1G，也会变成申请16G的显存，因此如果单机跑多进程TF，或者有其他进程也需要使用显存，可以将BFC allocator的显存分配策略适当修改一下，比如当显存占用超过一定阈值时，每次申请1G，而不是每次都翻倍，这样可以避免实际用8.1 而申请了16G显存的情况。
RoundedBytes 是将需要申请的显存字节数进行对齐，也就是实际返回的显存大小可能比申请的多，这样可以保证分配和回收不会产生大量的碎片(也可能有其他的好处)。上图可以看到首先尝试去拿到合适的chunk，如果找到了就直接返回指针，如果拿不到，说明已经申请的显存不够用了，就需要调用Extend 来申请更多的显存，然后再重新调用FindChunkptr 来找到合适的显存块。
上图是Extend函数的核心部分，主体思想就是申请翻倍的显存，如当前已经申请了2G， 就再调用sub_allocator的Alloc函数来申请2G显存，然后把当前申请了多少记录下来(curr_region_allocation_bytes_)，下次再翻倍。
二、获取当前实际使用显存量
有两个方法来知道
1、打log，我之前是修改了源码，每次调用AllocateRaw时，把实际显存的用量打出来，想打什么可以自己根据需要把DumpMemoryLog 函数稍微改一下，然后调用他就行了。这个方法稍微麻烦一些
2、sess run op tensorflow.contrib.memory_stats.MaxBytesInUse， 这个Op注册了CPU GPU 等一些device， 实际上也是去GPU allocator 里边拿的数据，直接把sess执行过程中占用最大的显存量返回来，这个用起来更加方便。
转载地址：http://tqhzb.baihongyu.com/