用沐神的方法阅读PyTorch FX论文

【GiantPandaCV导语】torch.fx对于PyTorch来说确实是一个比较好的工作，因为它消除了一些动态图和静态图的Gap。比如在图改写方面，torch.fx让PyTorch想做一些其它静态图框架的算子融合优化非常容易。并且torch.fx让后训练量化和感知训练量化以及AMP等的实现难度大大降低，这得益于我们可以直接在Python层操作这个IR，所以我认为这是一个不错的工作。尤其是对使用PyTorch开发的算法工程师来说，现在可以基于这个特性大开脑洞了。我之前围绕FX也做了一个QAT的工作，感兴趣可以阅读：基于OneFlow实现量化感知训练。torch.fx的卖点就是，它使用纯Python语言实现了一个可以捕获PyTorch程序的计算图并转化为一个IR的库，并且非常方便的在这个IR上做Pass，同时提供将变换后的IR Codegen合法的Python代码功能。我觉得算是达到了在Eager下写Pass就像做链表插入删除题目一样顺滑。

0x0. 动机

最近沐神在B站上分享了一些经典论文的阅读视频，我也跟着看了几个比如TransFormer，VIT等，很赞。所以我就想试试沐神这种论文阅读方法，找一篇论文来阅读一下。目前我比较关注的是工程方向的论文，正好上周PyTorch 放出了torch.fx的论文，所以我就以这篇论文为例来试试。沐神的论文阅读方法大概是这样（数字代表先后顺序）：

1. 标题
2. 摘要
3. 引言
4. 结论
5. 相关工作
6. FX特性
7. 实验
8. 评论

PyTorch FX论文的链接在：https://arxiv.org/pdf/2112.08429.pdf 。下面我就以沐神的论文阅读顺序来分享一下阅读体验，帮助大家搞清楚PyTorch FX这个特性到底是什么，以及它可以在PyTorch中发挥什么作用。

0x1. 标题

我们可以这样进行翻译：TORCH.FX：基于Python的深度学习实用程序捕获和变换。这个基于Python的深度学习实用程序这里就可以理解为基于PyTorch开发的模型程序了，然后这里的重点是捕获和变换。现在还不清楚这里的捕获和变换是什么意思，我们可以接着往下读。插一点言外话，我半年前就有关注FX，之前在OneFlow框架下也成功接入了FX并且做了一个QAT的任务，而那个时候FX还是没有论文的，所以我感觉这篇论文更像是一个对TORCH.FX这个特性的总结以及占位的那种感觉。

0x2. 摘要

摘要部分简单指明了像PyTorch这种基于动态图执行模式的深度学习框架虽然提升了用户的易用性。但是在一些真实场景中，用户可能是需要捕获和变换程序结构（也可以直接理解为神经网络的结构）来进行性能优化，可视化，分析和硬件调优等。为了解决这个痛点，PyTorch设计了torch.fx这个模块来做PyTorch程序的捕获和变换，这个模块是纯Python开发的。

这一节主要是讲了一下torch.fx的卖点，就是说动态图虽然易用性很强，但是图结构不能被提前感知和变换，但通过这篇论文的torch.fx模块，这件事就成啦！

0x3. 引言

早期的图模式或者叫define-and-run的静态图框架有Caffe，TensorFlow等，它们设计了一个表示图的IR，用户通过调用这些框架提供的API来构建IR。然后我们可以在这个IR上做程序微分，将IR切分到设备上实现并行，量化，性能优化等等。但这些事情一般都要求开发者在领域特定的语言上去做，比如以OneFlow的静态图模式为例，要做图切分，量化，性能优化等都需要基于C++去开发，并且调试也会相对比较难（要借助pdb，gdb等等专业工具）。

现在的eager模式或者叫define-