PyTorch 源码阅读笔记(7):TorchDynamo
关于 TorchDynamo torchdynamo 为 PyTorch 2.0 的新功能,可以在不修改代码的情况下,对大部分模型提速,基本的使用方式如下 import torch def fn(x, y): a = torch.cos(x).cuda() b = torch.sin(y).cuda() return a + b new_fn = torch.compile(fn, backend="inductor") input_tensor = torch.randn(10000).to(device="cuda:0") a = new_fn(input_tensor, input_tensor) TorchDynamo 原理 官方给出的 TorchDynamo 原理图如下 涉及到关于 Python 编译运行的内容参考 [python 编译运行过程](Python 代码编译运行过程(1):编译过程 | K’s blog (luokai.tech) 字节码优化 torchdynamo 通过捕捉 python 的 frame object 进行字节码优化,运行如下代码 from torch._dynamo import optimize import torch._dynamo.config import logging torch._dynamo.config.log_level = logging.INFO torch._dynamo.config.output_code = True @optimize() def toy_example(a, b): a *= 10 b = b + 1 return b for _ in range(100): toy_example(torch....