depyf:为机器学习研究人员打开 PyTorch 编译器的黑箱
机器学习
2024-03-22 v1 人工智能
编程语言
摘要
PyTorch \texttt{2.x} 引入了一个旨在加速深度学习程序的编译器。然而,对于机器学习研究人员而言,充分适应并发挥 PyTorch 编译器的潜力可能颇具挑战。该编译器在 Python 字节码层面运行,使其看起来像一个不透明的黑箱。为解决此问题,我们推出了 \texttt{depyf},一个旨在揭秘 PyTorch 编译器内部工作原理的工具。\texttt{depyf} 将 PyTorch 生成的字节码反编译回等效的源代码,并建立内存中代码对象与其磁盘上源代码对应物之间的连接。此功能使用户能够使用调试器逐行单步执行源代码,从而加深他们对底层过程的理解。值得注意的是,\texttt{depyf} 是非侵入式且用户友好的,其核心功能主要依赖于两个便捷的上下文管理器。该项目已\href{https://github.com/thuml/depyf}{公开可用},并被认定为\href{https://pytorch.org/ecosystem/}{PyTorch 生态系统项目}。
引用
@article{arxiv.2403.13839,
title = {depyf: Open the Opaque Box of PyTorch Compiler for Machine Learning Researchers},
author = {Kaichao You and Runsheng Bai and Meng Cao and Jianmin Wang and Ion Stoica and Mingsheng Long},
journal= {arXiv preprint arXiv:2403.13839},
year = {2024}
}
备注
16 pages, 2 figures