Torch.fx：面向 Python 中深度学习的实用程序捕获与变换
机器学习
2022-03-08 v2
摘要
现代深度学习框架提供嵌入于 Python 中的命令式、即时执行编程接口，以提供高效的开发体验。然而，深度学习从业者有时需要捕获并变换程序结构以进行性能优化、可视化、分析和硬件集成。我们研究了深度学习中用于程序捕获与变换的不同设计。通过为典型的深度学习用例而非长尾用例进行设计，可以创建一个更简单的程序捕获与变换框架。我们将此原则应用于 torch.fx——一个完全用 Python 编写、并为 ML 从业者高开发效率优化的 PyTorch 程序捕获与变换库。我们给出案例研究，展示 torch.fx 如何实现此前 PyTorch 生态系统中无法达成的工作流。
引用
@article{arxiv.2112.08429,
title = {Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python},
author = {James K. Reed and Zachary DeVito and Horace He and Ansley Ussery and Jason Ansel},
journal= {arXiv preprint arXiv:2112.08429},
year = {2022}
}
备注
14 pages, 8 figures, Accepted to MLSys 2022. v2: Added correctness information to evals, clarified 6.2.3, added transform runtime measurement