一文搞懂 TorchDynamo 原理

文章预览

【点击】 加入大模型技术交流群 原文：https://fkong.tech/posts/2023-05-20-dynamo/ 简介 PyTorch 2.0 的使命是更快、更 Pythonic 以及一如既往地支持动态特性。为了达到这个目的，PyTorch 2.0 引入了  torch.compile ，在解决 PyTorch 固有的性能问题的同时，把部分用 C++ 实现的东西引入 Python 中。PyTorch 2.0 利用了 4 个组件: TorchDynamo，AOTAutograd，PrimTorch 和 TorchInductor。本文以几个简单的案例讲解 TorchDynamo 的使用方法和实现原理。 PyTorch 2.0 TorchDynamo 的作用是从 PyTorch 应用中抓取计算图 ，相比于 TorchScript 和 TorchFX，TorchDynamo 更加灵活、可靠性更高。用过 TorchScript 的朋友知道，通过  jit.trace  或者  jit.script  把模型转化为 TorchScript 的过程困难重重，往往需要修改大量源代码。而 TorchFX 在捕获计算图时，遇到不支持的算子会直接报错，最常见的就是  if  语句。TorchDynamo 克服 ………………………………