随笔 - 128  文章 - 0  评论 - 75  阅读 - 11万

TorchDynamo 原理和示例

  TorchDynamo 是一个设计用于加速未修改的 PyTorch 程序的 Python 级即时(JIT)编译器。它通过 Python Frame Evaluation Hooks(Python 框架评估钩子)来实现这一目标,以便在运行时动态地生成和优化代码。这使得 TorchDynamo 可以有效地处理各种 Python 代码,包括包含控制流(如循环和条件语句)的代码,而无需进行任何修改。

一.TorchDynamo 工作流程

  TorchDynamo 的工作流程,如下所示:

  • 首先,TorchDynamo会捕获PyTorch应用中的计算图。
  • 然后,TorchDynamo将这些计算图转换为FX图,这是一种用于表示PyTorch程序的中间表示(IR)。
  • 最后,TorchInductor会消费这些FX图以产生优化的代码。

二.TorchDynamo 简单示例

  以下是一个简单的使用示例:

import torch
from torch import dynamo

# 定义一个简单的PyTorch模型
class Model(torch.nn.Module):
    def forward(selfx):
        return x * 2

model = Model()

# 使用TorchDynamo来优化model
optimized_model = dynamo.optimize(model)

# 现在可以使用优化后的model
x = torch.tensor([1.0])
y = optimized_model(x)

  在这个例子中,首先定义了一个简单的 PyTorch 模型,然后使用 TorchDynamo 的 optimize 函数来优化它。优化后的模型可以像普通的 PyTorch 模型一样使用。

参考文献

[1] torch.compiler_fine_grain_apis:https://pytorch.org/docs/2.1/torch.compiler_fine_grain_apis.html

[2] torchdynamo:https://github.com/pytorch/torchdynamo

[3] torch dynamo:https://github.com/pytorch/pytorch/tree/main/torch/_dynamo

[4] torch.compiler_deepdive:https://pytorch.org/docs/stable/torch.compiler_deepdive.html

[5] whats-behind-pytorch-2-0-torchdynamo-and-torchinductor-primarily-for-developers:https://pyimagesearch.com/2023/04/24/whats-behind-pytorch-2-0-torchdynamo-and-torchinductor-primarily-for-developers/




NLP工程化

1.本公众号以对话系统为中心,专注于Python/C++/CUDA、ML/DL/RL和NLP/KG/DS/LLM领域的技术分享。
2.本公众号Roadmap可查看飞书文档:https://z0yrmerhgi8.feishu.cn/wiki/Zpewwe2T2iCQfwkSyMOcgwdInhf

NLP工程化(公众号)

NLP工程化(星球号)

posted on   扫地升  阅读(515)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· 在鹅厂做java开发是什么体验
· 百万级群聊的设计实践
· WPF到Web的无缝过渡:英雄联盟客户端的OpenSilver迁移实战
· 永远不要相信用户的输入:从 SQL 注入攻防看输入验证的重要性
· 浏览器原生「磁吸」效果!Anchor Positioning 锚点定位神器解析
< 2025年2月 >
26 27 28 29 30 31 1
2 3 4 5 6 7 8
9 10 11 12 13 14 15
16 17 18 19 20 21 22
23 24 25 26 27 28 1
2 3 4 5 6 7 8

点击右上角即可分享
微信分享提示