我们之前的一篇文章，《优化 TensorFlow Lite 运行时内存》，讨论了 TFLite 的内存竞技场如何通过在张量之间共享缓冲区来最大限度地减少内存使用。这意味着我们可以在更小的边缘设备上运行模型。在今天的文章中，我将描述内存竞技场初始化的性能优化，以便我们的用户能够以很少的额外开销获得低内存使用的好处。

机器学习通常作为更大管道的一部分在设备上部署。TFLite 被使用是因为它快速且轻量级，但管道中的其他部分也必须快速。使用代表性数据在目标设备上进行分析使我们能够识别管道中最慢的部分，以便我们优化代码中最重要的部分。

在本文中，我将描述 TFLite 内存竞技场的分析和优化，并提供有关如何使用 Simpleperf 和可视化结果的说明。给出了示例命令。假设已安装 Android NDK，并且您拥有可以使用 adb 连接的开发设备。

Simpleperf

Simpleperf 附带一些脚本以使其更易于使用。run_simpleperf_on_device.py 将 simpleperf 推送到设备并使用给定的参数运行您的二进制文件。

/usr/lib/android-ndk/simpleperf/run_simpleperf_on_device.py record –call-graph fp /data/local/tmp/my_binary arg0 arg1 …

这将生成输出文件 perf.data，您必须将其复制回计算机。

adb pull /data/local/tmp/perf.data

然后您生成二进制缓存，其中包含以后生成有用配置文件所需的所有信息。

/usr/lib/android-ndk/simpleperf/binary_cache_builder.py -lib /your/binarys/folder -i perf.data

并生成用于可视化的协议缓冲区

/usr/lib/android-ndk/simpleperf/pprof_proto_generator.py --ndk_path=/path/to/android-ndk -i perf.data -o profile.proto

然后，您可以使用 pprof 显示这些结果

pprof -http :8888 profile.proto

然后在您的浏览器中打开 localhost:8888 以查看配置文件。我发现火焰图是最有用的

优化 TFLite 的内存竞技场

ArenaPlanner::ExecuteAllocations 占用了此模型运行时间的 54.3%。我原本以为会发现全连接层或卷积之类的机器学习算子是此模型的瓶颈，而不是运行时开销。这是一种特别糟糕的情况，内存竞技场开销并非每个模型都如此糟糕，但此处的改进将影响所有模型。此模型具有可变输入大小和许多动态张量，其输出大小直到算子评估后才能确定，这会触发频繁的张量重新分配。这确实是最糟糕的情况。让我们放大配置文件。

InterpreterInfo::num_tensors() 占用了运行时间的 10.4%。此函数如此昂贵的原因是因为它是一个虚拟函数，它调用另一个函数，并且在循环中被调用。我从来没想过会这样。

竞技场规划器不创建或销毁张量，因此张量数量是恒定的。让我们将其缓存起来。

下一个容易优化的部分是 InterpreterInfo::tensor(unsigned long)，这是一个虚拟函数，它进行边界检查并返回指向张量的指针。张量存储在一个数组中，因此让我们添加一个函数来获取指向该数组的指针。相应的代码提交记录可以在这里找到：[链接1](https://github.com/tensorflow/tensorflow/commit/7528df84ad0207ec88eb8324dee8e10cf79fda0d) 和 [链接2](https://github.com/tensorflow/tensorflow/commit/91cc89a28fcfeef7df3f257bd27086343d438610)。

经过这些简单的修改，该模型的运行时间减少了 25%，内存分配器的开销也减少了一半。Simpleperf 使得识别这些低效率问题变得容易！现在，让我们再次进行性能分析，以衡量这些修改的影响。

ArenaPlanner::CalculateAllocations 现在是开销最大的函数，占运行时间的 12.7%。该函数调用了两个函数：SimpleMemoryArena::Allocate 和 ArenaPlanner::CreateTensorAllocationVector。

虽然 ArenaPlanner::CreateTensorAllocationVector 的开销较小，但其代码更加简单，因此可能更容易优化。该函数识别了图中两个节点之间分配的张量，然后按大小对它们进行排序，因为它使用了贪婪按大小算法，该算法会先分配最大的张量。图的结构是恒定的，因此我们可以存储一个映射，用于存储在每个节点分配的张量。我们无需检查模型中的每个张量以查看它是否在两个节点之间分配，而可以通过遍历映射来识别要分配的张量。 ArenaPlanner::CreateTensorAllocationVector 的开销已从运行时间的 4.8% 降至 0.8%。相应的代码可以在这里找到：[链接3](https://github.com/tensorflow/tensorflow/commit/72c981ec392a4ec8b3e9fcae44c8be3217265437)。排序操作没有出现在性能分析中，因此我们忽略它。

接下来要关注的函数是 ArenaPlanner::ResolveTensorAllocation，它在进行之前的优化后占运行时间的 10.9%。该函数在分配后会重置每个张量的 `data` 指针。但是，这些指针并不总是改变。如果我们跟踪并仅更新那些改变的指针，会怎么样？经过这个 [修改](https://github.com/tensorflow/tensorflow/commit/c9b82169ab767a96aa1bf31556c186c9c91a9fe8) 后，ArenaPlanner::ResolveTensorAllocation 就不再出现在性能分析中。

现在让我们看一下分配和释放。 SimpleMemoryArena::Allocate 占运行时间的 7%，而 SimpleMemoryArena::Deallocate 占 6.8%。竞技场中所有分配的记录都存储在一个向量中，该向量按它们在内存竞技场中的偏移量排序。该排序数据结构中的条目会被插入、删除和搜索。这些操作在向量中都是 **O(N)** 的。使用 `std::multimap`（以偏移量作为键）会更好吗？我们需要使用 `multimap`，因为记录按偏移量排序，并且可能存在具有相同偏移量的多个张量。删除和插入是 **O(logN)**，但搜索仍然是 **O(N)**，因为我们正在搜索张量 ID，而不是偏移量。最好的方法是测试和分析。

用 `multimap` 替换向量实际上会降低竞技场代码的速度：它比使用向量慢了将近三倍！虽然这与直觉相悖，但这是在优化代码时常常见到的现象。集合或映射上的操作具有线性或对数复杂度，但是复杂度中也包含一个常数项。这个常数项的值高于向量的复杂度的常数项的值。我们还遍历了这些记录，对于向量来说，这比列表或 `multimap` 要便宜得多。我们还测试了列表，结果是它的速度是向量的两倍。

但是，释放仍然可以改进。 SimpleMemoryArena::Deallocate 会遍历这些记录，当它找到要释放的记录时，会将其从向量中删除。这具有 **O(N^2)** 的复杂度。性能分析中看到的 `memcpy` 来自对 `std::vector::erase` 的频繁调用。将记录标记为要删除，然后使用 `std::remove_if` 在一次遍历中删除它们会更高效得多。这里的第二个优化是关注张量是如何释放的：ArenaPlanner::ResetAllocationsAfter 会释放从节点到图末尾的所有张量。为了解决这个问题，添加了 SimpleMemoryArena::DeallocateAfter(int32_t node)，它会遍历所有记录，标记在节点之后分配的那些记录。 SimpleMemoryArena::ResolveDeallocations 会在一次遍历中删除这些记录，从而使释放操作变为 **O(N)**。经过这些修改，ResetAllocationsAfter 就不再出现在性能分析中！相应的代码提交记录可以在这里找到：[链接4](https://github.com/tensorflow/tensorflow/commit/509b811aba38294f451dd7beea3c27558ce1f7da) 和 [链接5](https://github.com/tensorflow/tensorflow/commit/9e582c01e1a6f813d4c76a845fe3802fbc3140f8)。

性能分析现在看起来完全不同，张量分配的开销已从运行时间的 49.9% 降至 11%。这个性能分析看起来合理得多。 SimpleMemoryArena::Allocate 是最后一个需要优化的函数。对于每个张量，该函数会遍历记录向量，试图找到当前分配的空间。这具有 **O(N^2)** 的复杂度。这是贪婪按大小算法的根本局限性。高效利用内存是以增加开销为代价的。虽然复杂度无法降低，但 **N** 可以降低。我们按执行顺序处理节点。对于已在已执行的节点上释放的张量的分配信息不再需要，它只会拖慢速度。记录会定期清除，以便仅考虑活动记录。在大型模型上，这会显著降低 **N**。 ArenaPlanner::ExecuteAllocations 现在不再是最昂贵的函数！它已从运行时间的 11% 降至 6%，而全连接算子现在是性能分析中开销最大的函数，这正是我们在分析神经网络推理时所期望的。

这正是神经网络性能分析应该呈现的样子。应该将时间花在运行模型的算子上，而不是在推理运行时上。

优化的内存竞技场现在作为 TensorFlow 2.13 的一部分公开提供。

下一步

今天的文章向您展示了 Simpleperf 如何帮助我们找到 TFLite 内存竞技场中易于修复的低效率问题，而这些问题仅仅通过查看代码是无法发现的。Pprof 可以显示带注释的源代码、反汇编和图形，使您能够轻松找到设备上管道中的瓶颈。