北京建设官方网站,网易联合创新中心,校园网站建设合同百度文库,用前端做的比较酷的网站背景
内存管理是AI计算中非常重要的一部分。我们希望模型计算时占用内存尽可能小#xff0c;这样我们训练或推理时就可以用更大的batch size使其尽快收敛#xff0c;或者提高吞吐率。又或者让我们可以使用参数更多、或更复杂的模型从而达到更好的准确率。由于现代深度学习模…背景
内存管理是AI计算中非常重要的一部分。我们希望模型计算时占用内存尽可能小这样我们训练或推理时就可以用更大的batch size使其尽快收敛或者提高吞吐率。又或者让我们可以使用参数更多、或更复杂的模型从而达到更好的准确率。由于现代深度学习模型大多在GPU上运行而GPU的显存相比CPU小得多因此我们这里主要关注的是GPU memory。
首先看下我们需要重点关注哪些GPU memory。对于模型在计算中需要用到的GPU memory论文《Estimating GPU Memory Consumption of Deep Learning Models》做了比较具体的总结。对推理计算而言主要有这么几类
权重参数Weight Parameter这个不用多说模型中的参数。中间张量Intermediate Tensor网络中每层的输出与输入张量。其它计算中需要的临时内存如kernel中使用的一些memorycuDNN的workspace还有一些常驻的内存如CUDA context。
其中第三类本身占的空间不大也比较难优化。第一类减少权重内存占用的话可以通过一些模型压缩方法如量化剪枝。之前写过一些相关文章如《闲话模型压缩之量化Quantization篇》和《闲话模型压缩之网络剪枝Network Pruning篇》有兴趣可以参考。相比之下第二类即中间张量的优化空间更大因此很多业界的工作也是针对它来优化。优化的思路有很多比如
内存重用Memory reuse由于有些中间张量的生命周期间互不重叠因此可以reuse。MegEngine, IREE, TensorRT等框架都做了memory usage相关的优化。重计算Recomputation该技术主要用于训练中。它将模型中的一些节点作为checkpoint其它节点的输出可丢弃当在计算梯度时需要时通过最近的checkpoint重新计算生成。因此被称为checkpointing技术。该问题也被称为tensor rematerialization优化问题。相关论文如适用于静态网络的offline方法《Training Deep Nets with Sublinear Memory Cost》适用于动态网络的online方法《Dynamic Tensor Rematerialization》建模为MILP进行求解的《Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization》等。交换Swap基本思想是将显存中的数据交换到CPU上相当于把GPU memory当成CPU memory的cache。一些塞不进GPU显存的层如DLRM模型的embedding层可能会用到这种技术。相关的论文如《Supporting Massive DLRM Inference Through Software Defined Memory》《vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design》等。压缩Compression将数据进行压缩如《Gist: Efficient Data Encoding for Deep Neural Network Training》根据特定层输出特点对层的输出即feature map数据进行编码压缩。…
后面几种都会一定程度上牺牲性能这里我们主要关注第一种。它在对延迟关注的推理场景用得尤其多一些。比如TensorFlow Lite利用该技术可以显著减少内存占用详见https://blog.tensorflow.org/2020/10/optimizing-tensorflow-lite-runtime.html。
对于网络前面层的输出到计算后面的层时可能已经不再使用了。换句话说对于中间层的输出张量它们的生命周期可能是没有重叠的对于它们便可以进行重用。下面是最简单情况下的示意图 其中Task 1写数据到Tensor 1交给Task 2Task 2处理后将结果写入Tensor 2交给Task 3。因为Tensor 1与Tensor 2生命周期并不重叠所以它们可以重用同一个Tensor。典型的如神经网络中的一些element wise操作。
但实际中的情况远没有这么简单。网络中不总是线性结构另外张量的大小可能各不相同给重用带来困难。要得到其最优的分配策略是一个NP-complete问题。因此实际当中我们可以倾向于一些heuristic方法这样可以在合理时间内得到一个近似最优解。
那如何优化呢在不少地方如TensorRT会提到基于register allocation的思想。任何一本编译器的教材中都会介绍register allocation在此不展开。大体会先通过liveness analysis得到变量的live range然后根据它生成inference graph转为着色问题来解。业界也有采用这种方法的做法如《Memory Allocation for Neural Networks using Graph Coloring》。但是memory planning与register allocation所面临的问题还是有所区别的比如
大小不确定Register的大小基本是相同或者说是差不多的而memory的大小可能差异很大重用一块过大或过小的memory会产生问题。拷贝成本高Register拷贝一下还好但memory拷贝开销比较大尤其是大块的memory。因此理想情况下我们希望不要拷贝。Fallback机制Register实在分配不了会产生spill即放到memory中。虽然memory要不不够理论上也能往更下一层存储器上搬业界也有这方面的研究但很多情况下因为时延等原因不会这么做。
另外从静态/动态角度内存的管理方式大体有动态分配与静态规划两种
动态分配Dynamic Memory Allocation内存分配在运行时进行。由于从系统中分配的开销较大通常维护一个memory pool。需要时从中分配不再需要时放回到pool。如TensorFlow中的BFC allocator。静态规划Static Memory Planning内存分配在运行前进行常见于基于编译器的计算框架。通过规划进一步减少内存使用减少OOM带来的不确定性同时最少化运行时内存管理的开销。如MXNet与MegEngine/MegCC中的static memory planning。
光看概念有些抽象下面就以一个实例 - TensorFlow LiteTF Lite中的内存优化来看看具体的实现。其实在论文《Efficient Memory Management for Deep Neural Net Inference》与《On-Device Neural Net Inference with Mobile GPUs》中对其原理已经介绍得比较清楚了。下面主要是结合代码理解下它的实现。
代码走读
基础数据结构
先来看几个关键数据结构。它们定义在types.h文件中。结构体TensorUsageRecord即论文中提到的Tensor usage record用于记录张量的使用记录。它主要包含三个信息tensor size, 以及第一个与最后一个使用它的task。代表这两个task的成员first_task与last_task即它的生命周期。
using UsageGraph std::vectorstd::vectorsize_t;template typename TensorSizeT
struct TensorUsageRecord {TensorSizeT tensor_size;TaskId first_task;TaskId last_task;...
};注意它是个模板类有针对size_tuint2uint3与BHWC的特化实现在memory_management.c文件。
结构体ObjectsAssignment与OffsetsAssignment都用于存放分配结果。
// Information about assignment of tensors to shared objects
template typename TensorSizeT
struct ObjectsAssignment {// shared_object_ids_[i] is ID of shared object, that tensor i will be using.std::vectorsize_t object_ids;// shared_object_sizes_[i] is a size of shared object with ID equal to i.std::vectorTensorSizeT object_sizes;
};// Information about assignment of tensors to offsets for the case, when all of
// them are going to be allocated in one continuous memory block.
struct OffsetsAssignment {std::vectorsize_t offsets;size_t total_size;
};它们对应后面会提到的两种分配方式。前者用于shared object指可以用于多个tensor的内存块分配后者用于从大块连续内存中分配子内存区域。
为了解它的使用可以参考它的测试用例memory_management_test.cc。比较典型的有几个caseOneRecordChainRecordsComplexRecords分别对于只有一个节点链式即线性计算图和复杂计算图。
以ChainRecords这个case为例
TEST(Model, ChainRecords) { std::vectorTensorUsageRecordsize_t usage_records{ {/*size*/16, /*first*/0, /*last*/1}, {/*size*/8, /*first*/1, /*last*/2}, {/*size*/64, /*first*/2, /*last*/3}, {/*size*/32, /*first*/3, /*last*/4}, {/*size*/8, /*first*/4, /*last*/5}, }; ObjectsAssignmentsize_t assignment; ASSERT_TRUE( AssignObjectsToTensors(usage_records, MemoryStrategy::NAIVE, assignment) .ok()); EXPECT_THAT(assignment.object_ids, ElementsAre(0, 1, 2, 3, 4)); EXPECT_THAT(assignment.object_sizes, ElementsAre(16, 8, 64, 32, 8)); ...可以看到其中最核心的是AssignObjectsToTensors()函数。该函数基于由TensorUsageRecord数组表示的张量使用记录按拓扑序排列根据指定的分配策略由MemoryStrategy表示计算得到分配结果由ObjectsAssignment对象表示。
Object分配方式
注意AssignObjectsToTensors()是个模板函数根据TensorUsageRecord的类型不同有多种实现。以最简单的TensorUsageRecordsize的情况即tensor的大小以一个size_t类型表示为例相关代码如下
template
absl::Status AssignObjectsToTensors(const std::vectorTensorUsageRecordsize_t usage_records,MemoryStrategy strategy, ObjectsAssignmentsize_t* assignment,const UsageGraph* reallocation_graph) {switch (strategy) {case MemoryStrategy::NAIVE:return NaiveAssignment(usage_records, assignment);case MemoryStrategy::EQUALITY:return EqualityAssignmentWithHash(usage_records, assignment);case MemoryStrategy::GREEDY_IN_ORDER:return GreedyInOrderAssignment(usage_records, assignment,reallocation_graph);case MemoryStrategy::GREEDY_BY_BREADTH:return GreedyByBreadthAssignment(usage_records, assignment);case MemoryStrategy::GREEDY_BY_SIZE:return GreedyBySizeDistPriorityAssignment(usage_records, assignment);case MemoryStrategy::GREEDY_BEST:return BestGreedy(usage_records, assignment);case MemoryStrategy::MINCOSTFLOW:return MinCostFlowAssignment(usage_records, assignment);default:return absl::InternalError(MemoryStrategy is not supported with current tensor size type.);}return absl::OkStatus();
}可以看到它的主体部分就是根据指定策略调用相应的分配算法。这几种策略分别是
NATIVE
NaiveAssignment函数将每个张量分配独立的memory object。实现在native_assignment.h文件中。该算法为每个张量分配一个新的shared object。这是最简单也是最浪费内存的做法。代码逻辑比较好理解不多说了。
EQUALITY
EqualityAssignmentWithHash()函数实现于equality_assignment.h文件中。它适用于TensorSizeT为hashable type的情况unhashable type的情况使用EqualityAssignment()函数。
该算法维护两个关键数据结构一个是pool它是一个map。其键值为size值为目前free即空闲的且size与键值相同的shared objects。另一个是优先队列objects_in_use它保存目前在使用的share objects并按size排序。整个过程遍历所有的tensor usage record对于每个tensor usage record先将队列objects_in_use中相对于当前tensor已不再使用的shared object弹出放入pool。然后在pool中找有没有与当前tensor的size匹配的shared object如有就重用没有的话就新创建一个shared object。过程示意图如下 注意这里只是做memory planning所以不用真正做内存分配。
GREEDY_IN_ORDER
函数GreedyInOrderAssignment()实现在greedy_in_order_assignment.h文件中。它维护与前一算法中相似的两个数据结构一个是存放free shared objects的pool以object size排序。另一个是存放正在使用的shared object的优先队列objects_in_use以last_task排序。该算法主体部分遍历tensor usage records。在第一步中首先看哪些object不再使用将它们放入pool。这一步与前面算法一样。
然后尝试将pool中的shared object分配给当前tensor。前面是需要严格匹配才能重用这里放宽了一些允许在size不严格一致时也能重用。这里在从pool中找可用的shared object时不是用的find()函数而是用的是二分查找lower_bound()即查找不小于当前tensor的size的第一个元素。这样做保证找到的shared object如有能容纳当前tensor同时又使浪费的空间尽可能小。然后尝试检查前一个元素其shared object size与tensor size的差值是否比前面找到的元素更小。如果是就选它了。这样使得对shared object的size改变尽可能得小就能满足当前tensor的需求。比如pool中有size分别为1, 3, 6的shared object而当前tensor为5。这种情况下就会找到size为3这个shared object因为它与5的差值是最小的。可以看到它每一步尽可能找与当前tensor的size尽可能接近的shared object进行重用体现了贪心的思想。
当前面的步骤中找到合适的shared object就把它从pool中拿走将之分配给当前tensor分配信息写入assignment。同时将该信息记录在objects_in_use中。如果shared object的size小于tensor会增大shared object的size。如果没有找到合适的shared object则创建新的shared object。
GREEDY_BY_BREADTH
函数GreedyByBreadthAssignment()实现在greedy_by_breadth_assignment.cc文件中。与前面的贪心算法类似主要区别在于它优先考虑breadth大的task。Breadth表示该task执行时所有tensor的size之和记录于TaskProfile对象中。TaskProfile是一个vector元素表示task执行时还在使用的张量以size降序排序。obj_schedules是SharedObjectSchedule的vector。SharedObjectSchedule记录了shared object对应的所有tensor usage record。
// Set of usage records for all tensors assigned to the shared object, ordered
// by first_task.
using SharedObjectSchedule std::setTensorUsageRecordsize_t; struct TaskBreadthWithId {size_t breadth;TaskId task_id;TaskBreadthWithId(size_t breadth, size_t task_id) : breadth(breadth), task_id(task_id) {} // Default order of TaskBreadthWithId is increasing order of their breadth. bool operator(const TaskBreadthWithId other) const {return breadth other.breadth;}
}; 算法首先调用CalculateTaskProfiles()函数计算出所有task的TaskProfile它包含了task的breadth信息。然后以breadth非递增顺序遍历所有task。对于每个task遍历其所有的tensor为它们分配shared object。分配过程中考虑obj_schedules中的shared object是否可重用。遍历其中元素如果不合适就跳过否则用lower_bound()函数找到其中不小于当前所需size的第一个shared object。如果找到的话检查一下如果合法就将该shared object作为最优候选。如果没找到就创建新的shared object。
GREEDY_BY_SIZE
函数GreedyBySizeDistPriorityAssignment()实现在greedy_by_size_assignment.cc文件中。它的流程大体与前一种相似区别在于优先考虑哪些tensor时所用的策略。该方法使用了一种更加复杂的heuristic。核心数据结构是SizeDistPriorityInfo
struct SizeDistPriorityInfo {// - Tensor with leftmost position in positional maximums vector has higher// priority;// - If two tensors have equal position, the one, that has usage interval with// smallest positive distance (best_dist) to some of already assigned tensors,// has higher priority;// - If two tensors have equal position and best_dist, the one with greater// tensor_size has higher priority.bool operator(const SizeDistPriorityInfo other) const {return position other.position ||(position other.position (best_dist other.best_dist || (best_dist other.best_dist tensor_size other.tensor_size)));}// Recalculate best distance and best object, based on precalculated distances// in vector dist.void RecalcBestDist() {best_dist kNotAssigned;for (size_t obj_id 0; obj_id dist.size(); obj_id) {if (dist[obj_id] best_dist) {best_dist dist[obj_id];best_object obj_id;}}}size_t position;size_t tensor_size;std::vectorsize_t dist;size_t best_dist;size_t best_object;size_t tensor_usage_id;
};根据上面的注释在positional maximum vector中位置越靠左的tensor优先级越高。如果在该vector中位置一样则与已分配的tensor的positive distance小的优先级更高。如果前两个标准还决不出高下那size较大的tensor优先级更高。这个优先级考虑了tensor的size也考虑了与时序上的局部性。
主流程中首先调用CalculatePositionalMaximums()函数计算positional maximums vector。即每个shared object size的lower bound的数组。然后根据这个信息填好SizeDistPriorityInfo数组。接下来分两层循环。它会以SizeDistPriority递增顺序遍历所有tensor。即对于每个tensor按前面计算得到的优先级信息找到优先级最高的tensor以及相应的object。如果找到就按之进行分配没找到就创建新的shared object。分配完成后SizeDistPriority信息会产生变化因此需要进行相应的更新。
GREEDY_BEST
函数BestGreedy()实现在memory_management.cc文件中。它结合了前两种贪心算法。结合方式比较直观即把GREEDY_BY_SIZE与GREEDY_BY_BREADTH都跑下看哪个的总用量少即更优就选哪个。 RETURN_IF_ERROR(GreedyBySizeDistPriorityAssignment(usage_records, assignment));ObjectsAssignmentsize_t assignment_by_breadth;if (GreedyByBreadthAssignment(usage_records, assignment_by_breadth).ok() TotalSize(assignment_by_breadth) TotalSize(*assignment)) {std::swap(*assignment, assignment_by_breadth);}MINCOSTFLOW
函数MinCostFlowAssignment()实现在min_cost_flow_assignment.cc文件中。 它使用Minimum-cost flow matching algorithm。它将问题建模为一个minimum-cost flow problemMCFP。对于原问题它创建一个auxiliary flow graph对于每个intermediate tensor在图中插入两个对应节点另外创建两个特殊节点-source与sink。然后按一定规则向图中添加有向边具体可参见论文《On-Devie Neural Net Inference with Mobile GPUs》。创建完后使用Shortest Path Faster AlgorithmSPFA求解。
TF Lite中虽然提供了多种机制但至于哪种好不一定所以实际使用当中可能要都试一下。
Offset分配方式
除了这种以shared object分配的模式外TF Lite还支持一种在一整块内存块分配的模式。两种模式区别示意图 代码中检测如果支持sub buffer方式的话会试图以memory中的切块以offset表示进行分配。一般使用在先分配一大块连续memory然后在里边“切”小块的情况。这种情况下会调用AssignOffsetsToTensors()函数继而调用GreedyBySizeAssignment()函数。这是论文中提到的Offet Calculation方法。它可以看作是二维的strip packing problem。在主逻辑函数GreedyBySizeAssignment()中需要注意的是两个比较关键的数据结构ordered_records是按size排序后的TensorUsageRecord。ordered_allocs是已经分配好memory的tensor按offset排序。主要逻辑包含两层循环。外循环按序遍历TensorUsageRecord然后内循环中遍历已分配的内存块。如果生命周期不重叠则会考察它前面空出的区域是否能容纳下当前的tensor。如果能容纳且浪费的区域最小则将当前tensor塞进去。找一圈都找不到的话就在最右边新分配一段。
框架对接
看了内存分配的实现再看下它是如何结合进框架用于模型的计算的。我们知道TF Lite会将计算放到相应的后端如HexagonOpenGL, OpenCL, Metal等。这些后端实现放在delegates目录下。以GPU的OpenCL后端为例相应的实现主要位于tensorflow/tensorflow/lite/delegates/gpu/cl/inference_context.cc文件中。函数InitFromGpuModel()调用AllocateMemory()函数函数AllocateMemory继而调用AllocateBufferBasedTensors()函数或者AllocateStrongShapesTensors函数。前者用于buffer based的tensor参见IsBufferBased()函数比如clCreateBuffer()分配的内存。其它的则用AllocateStrongShapesTensors()函数进行处理。函数AllocateStrongShapesTensors()用于那些具有特定layout的tensor参考delegates/gpu/common/shape.h调用AssignObjectsToTensors()函数时传入的分配策略是MemoryStrategy::EQUALIT。
下面看下AllocateBufferBasedTensors()函数。其中核心函数GetBufferAsignment()负责buffer的分配。它先调用GetUsages()函数得到网络中张量的使用信息然后基于它构成TensorUsageRecord数组。有了这些信息就可以调用前面提到的AssignObjectsToTensors()函数进行内存对象的分配了。默认用的分配策略是MemoryStrategy::GREEDY_BEST。
另外如果支持sub buffer的话还会调用AssignOffsetsToTensors()进行分配。最后判断与前面以buffer为单位AssignObjectsToTensors()函数的分配方式得到的结果哪种好即使用的memory少从而选择更优的一种。
最后总结一下整个调用流程作为收尾吧
