利用XNNPack权重缓存进行内存高效推理

635 阅读4分钟

发布者:Zhi An Ng和Marat Dukhan,谷歌

XNNPack是默认的TensorFlow Lite CPU推理引擎,用于浮点模型,并在移动、桌面和网络平台上提供有意义的速度提升。XNNPack采用的优化之一是将卷积、深度卷积、转置卷积和全连接运算符的静态权重重新打包到一个为推理计算优化的内部布局。在推理过程中,重新打包的权重以一种对处理器管道友好的顺序模式被访问。

推理延迟的减少是有代价的:重新打包基本上在XNNPack内创建了一个额外的权重副本。当TensorFlow Lite模型被内存映射时,操作系统最终会释放权重的原始副本并使开销消失。然而,一些用例需要为同一个模型创建多个TensorFlow Lite解释器的副本,每个都有自己的XNNPack委托。由于属于不同的TensorFlow Lite解释器的XNNPack委托是不知道对方的,每一个都会创建自己的重新打包的权重副本,内存开销随着委托实例的数量而线性增长。此外,由于模型中的原始权重是静态的,XNNPack中的重新打包的权重在所有实例中也是一样的,因此这些副本是浪费和不必要的。

权重缓存是一种机制,它允许XNNPack委托的多个实例加速同一个模型,以优化它们对重新打包的权重的内存使用。有了权重缓存,所有的实例都使用相同的底层重新打包的权重,从而使内存使用量保持不变,无论创建多少个解释器实例。此外,由于权重缓存消除了重复的权重,可以通过提高处理器的缓存层次的效率来改善性能。注意:权重缓存是一个只通过C++ API提供的选择功能。

下面的图表显示了创建多个实例(横轴)的高水位内存使用量(纵轴)。它比较了不使用权重缓存的基线和使用权重缓存的软最终化。使用权重缓存时的峰值内存使用量相对于创建的实例数量来说增长得更慢。在这个例子中,使用权重缓存可以使你在相同的峰值内存预算下创建的实例数量增加一倍。

weights cache对象由TfLiteXNNPackDelegateWeightsCacheCreate 函数创建,并通过委托选项传递给XNNPack委托。然后,XNNPack委托将使用权重缓存来存储重新打包的权重。重要的是,权重缓存必须在任何推理调用之前被最终确定。

// Example demonstrating how to create and finalize a weights cache.

有两种方法来最终确定权重缓存,在上面的例子中,我们使用TfLiteXNNPackDelegateWeightsCacheFinalizeHard ,它执行最终确定。终结的内存开销最小,因为它将把权重缓存使用的内存削减到绝对最小。然而,在硬终结之后,不能用这个权重缓存对象创建新的委托--使用这个缓存的XNNPack委托实例的数量是事先固定的。另一种最终处理是最终处理。软终结有更高的内存开销,因为它在权重缓存中留有足够的空间用于一些内部记账。软终结的优点是相同的权重缓存可以用来创建新的XNNPack委托实例,只要这些委托实例使用完全相同的模型。如果委托实例的数量不是固定的或事先知道的,这就很有用。

// Example demonstrating soft finalization and creating multiple

接下来的步骤

有了权重缓存,使用XNNPack进行批量推理将减少内存使用,从而获得更好的性能。在README中阅读更多关于如何使用XNNPack的权重缓存,并在XNNPack的GitHub页面报告任何问题。

为了保持最新的信息,你可以阅读TensorFlow博客,关注twitter.com/tensorflow,或者订阅youtube.com/tensorflow。如果你已经建立了一些你想分享的东西,请提交给我们的社区焦点,地址是goo.gle/TFCS。如需反馈,请在GitHub上提交问题或在TensorFlow论坛上发帖。谢谢你