一句话:提出统一插桩框架 Amanda,把 DNN 的追踪、剪枝等分析优化任务抽象为算子级插桩,兼容多后端与多执行模式。
深度神经网络的分析(如 tracing)与优化(如剪枝)任务在 TensorFlow/PyTorch 等后端中各有特定需求和临时实现,开发者需应对碎片化接口并为不同模型改写代码。论文将这些任务的实现形式化为「神经网络插桩」,即在 DNN 的算子层引入插桩,从而把分析优化任务抽象为可作用于各类模型的插桩工具。Amanda 设计了两级 API:用户级 API 提供统一的算子粒度插桩接口,屏蔽不同后端差异;内部则采用以回调为中心的 API,管理并优化原始代码与插桩代码在不同后端中的执行。该设计使框架能覆盖 tracing、profiling、剪枝、量化等多种用例,跨越 TensorFlow/PyTorch 等后端以及图模式与 eager 模式,且性能开销通常控制在 5% 以内。
The success of deep neural networks (DNNs) has sparked efforts to analyze (e.g., tracing) and optimize (e.g., pruning) them. These tasks have specific requirements and ad-hoc implementations in current execution backends like TensorFlow/PyTorch, which require developers to manage fragmented interfaces and adapt their codes to diverse models. In this study, we propose a new framework called Amanda to streamline the development of these tasks. We formalize the implementation of these tasks as neural network instrumentation, which involves introducing instrumentation into the operator level of DNNs. This allows us to abstract DNN analysis and optimization tasks as instrumentation tools on various DNN models. We build Amanda with two levels of APIs to achieve a unified, extensible, and efficient instrumentation design. The user-level API provides a unified operator-grained instrumentation API for different backends. Meanwhile, internally, we design a set of callback-centric APIs for managing and optimizing the execution of original and instrumentation codes in different backends. Through these design principles, the Amanda framework can accommodate a broad spectrum of use cases, such as tracing, profiling, pruning, and quantization, across different backends (e.g., TensorFlow/PyTorch) and execution modes (graph/eager mode). Moreover, our efficient execution management ensures that the performance overhead is typically kept within 5%.
@inproceedings{GuanQL0YLF0Z00L24,
title = {Amanda: Unified Instrumentation Framework for Deep Neural Networks},
author = {Yue Guan and Yuxian Qiu and Jingwen Leng and Fan Yang and Shuo Yu and Yunxin Liu and Yu Feng and Yuhao Zhu and Lidong Zhou and Yun Liang and Chen Zhang and Chao Li and Minyi Guo},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3617232.3624864}
}