评价此页

★ ★ ★ ★ ★

intermediate/torch_compile_tutorial

在 Google Colab 中运行

注意

跳转至页面底部下载完整示例代码。

`torch.compile` 简介#

创建日期：2023年3月15日 | 最后更新：2026年4月1日 | 最后验证：2024年11月5日

作者： William Wen

torch.compile 是加速 PyTorch 代码的新方式！torch.compile 通过将 PyTorch 代码实时（JIT）编译为优化后的内核，从而使代码运行得更快，同时只需极少的代码修改。

torch.compile 通过追踪您的 Python 代码并查找其中的 PyTorch 操作来实现这一目标。难以追踪的代码会导致图中断（graph break），这意味着优化机会的丢失，而不是产生错误或静默错误。

torch.compile 在 PyTorch 2.0 及更高版本中可用。

本简介涵盖了 torch.compile 的基本用法，并展示了 torch.compile 相较于我们之前的 PyTorch 编译器解决方案 TorchScript 的优势。

如需查看关于真实模型的端到端示例，请查阅我们的 torch.compile 端到端教程。

如需排查问题并深入了解如何将 torch.compile 应用于您的代码，请查阅 torch.compile 编程模型。

内容

本教程所需的 pip 依赖项

torch >= 2.0
numpy
scipy

系统要求 - C++ 编译器，例如 g++ - Python 开发包 (python-devel/python-dev)

基本用法 #

在本教程中，我们开启了一些日志记录，以帮助了解 torch.compile 在底层做了什么。以下代码将打印出 torch.compile 追踪到的 PyTorch 操作。

import torch

torch._logging.set_logs(graph_code=True)

torch.compile 是一个装饰器，可以接受任意 Python 函数。

def foo(x, y):
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b


opt_foo1 = torch.compile(foo)
print(opt_foo1(torch.randn(3, 3), torch.randn(3, 3)))


@torch.compile
def opt_foo2(x, y):
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b


print(opt_foo2(torch.randn(3, 3), torch.randn(3, 3)))

TRACED GRAPH
 ===== __compiled_fn_1_2dca5809_c058_4275_a15b_63f190b006d5 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[3, 3][3, 1]cpu", L_y_: "f32[3, 3][3, 1]cpu"):
        l_x_ = L_x_
        l_y_ = L_y_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:74 in foo, code: a = torch.sin(x)
        a: "f32[3, 3][3, 1]cpu" = torch.sin(l_x_);  l_x_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:75 in foo, code: b = torch.cos(y)
        b: "f32[3, 3][3, 1]cpu" = torch.cos(l_y_);  l_y_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:76 in foo, code: return a + b
        add: "f32[3, 3][3, 1]cpu" = a + b;  a = b = None
        return (add,)


tensor([[ 1.0236,  1.0035,  0.1669],
        [ 1.2657,  1.0057, -0.0585],
        [-0.5140,  0.9709, -0.3063]])
TRACED GRAPH
 ===== __compiled_fn_3_ec0c3f1c_c73b_4109_ba59_7d8a969b6026 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[3, 3][3, 1]cpu", L_y_: "f32[3, 3][3, 1]cpu"):
        l_x_ = L_x_
        l_y_ = L_y_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:85 in opt_foo2, code: a = torch.sin(x)
        a: "f32[3, 3][3, 1]cpu" = torch.sin(l_x_);  l_x_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:86 in opt_foo2, code: b = torch.cos(y)
        b: "f32[3, 3][3, 1]cpu" = torch.cos(l_y_);  l_y_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:87 in opt_foo2, code: return a + b
        add: "f32[3, 3][3, 1]cpu" = a + b;  a = b = None
        return (add,)


tensor([[ 0.4807,  0.2284, -1.2120],
        [ 0.2812,  0.2407, -0.0167],
        [ 1.1850, -0.1689,  0.1028]])

torch.compile 会递归应用，因此顶级已编译函数内部嵌套的函数调用也将被编译。

def inner(x):
    return torch.sin(x)


@torch.compile
def outer(x, y):
    a = inner(x)
    b = torch.cos(y)
    return a + b


print(outer(torch.randn(3, 3), torch.randn(3, 3)))

TRACED GRAPH
 ===== __compiled_fn_5_d2dfbd45_fc55_4528_a29e_e292e422d8c7 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[3, 3][3, 1]cpu", L_y_: "f32[3, 3][3, 1]cpu"):
        l_x_ = L_x_
        l_y_ = L_y_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:98 in inner, code: return torch.sin(x)
        a: "f32[3, 3][3, 1]cpu" = torch.sin(l_x_);  l_x_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:104 in outer, code: b = torch.cos(y)
        b: "f32[3, 3][3, 1]cpu" = torch.cos(l_y_);  l_y_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:105 in outer, code: return a + b
        add: "f32[3, 3][3, 1]cpu" = a + b;  a = b = None
        return (add,)


tensor([[ 1.3105,  0.1338, -0.3341],
        [ 1.0808,  0.3435,  0.8590],
        [ 0.4188,  0.3586,  1.8927]])

我们还可以通过调用 torch.nn.Module 实例的 .compile() 方法，或者直接对模块使用 torch.compile 来优化它。这等同于对模块的 __call__ 方法进行 torch.compile 处理（它会间接调用 forward）。

t = torch.randn(10, 100)


class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = torch.nn.Linear(3, 3)

    def forward(self, x):
        return torch.nn.functional.relu(self.lin(x))


mod1 = MyModule()
mod1.compile()
print(mod1(torch.randn(3, 3)))

mod2 = MyModule()
mod2 = torch.compile(mod2)
print(mod2(torch.randn(3, 3)))

TRACED GRAPH
 ===== __compiled_fn_7_fabf108e_9288_4c11_b7f5_f296688d238f =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_self_modules_lin_parameters_weight_: "f32[3, 3][3, 1]cpu", L_self_modules_lin_parameters_bias_: "f32[3][1]cpu", L_x_: "f32[3, 3][3, 1]cpu"):
        l_self_modules_lin_parameters_weight_ = L_self_modules_lin_parameters_weight_
        l_self_modules_lin_parameters_bias_ = L_self_modules_lin_parameters_bias_
        l_x_ = L_x_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:126 in forward, code: return torch.nn.functional.relu(self.lin(x))
        linear: "f32[3, 3][3, 1]cpu" = torch._C._nn.linear(l_x_, l_self_modules_lin_parameters_weight_, l_self_modules_lin_parameters_bias_);  l_x_ = l_self_modules_lin_parameters_weight_ = l_self_modules_lin_parameters_bias_ = None
        relu: "f32[3, 3][3, 1]cpu" = torch.nn.functional.relu(linear);  linear = None
        return (relu,)


tensor([[0.0509, 0.0000, 0.9762],
        [0.0000, 0.0000, 0.3544],
        [0.0000, 0.0000, 0.1016]], grad_fn=<CompiledFunctionBackward>)
tensor([[0.0000, 0.0000, 0.3575],
        [0.0741, 0.0000, 0.1254],
        [0.0000, 0.0000, 0.2622]], grad_fn=<CompiledFunctionBackward>)

演示加速效果 #

现在让我们演示 torch.compile 如何加速一个简单的 PyTorch 示例。如需在更复杂的模型上进行演示，请参阅我们的 torch.compile 端到端教程。

def foo3(x):
    y = x + 1
    z = torch.nn.functional.relu(y)
    u = z * 2
    return u


opt_foo3 = torch.compile(foo3)


# Returns the result of running `fn()` and the time it took for `fn()` to run,
# in seconds. We use CUDA events and synchronization for the most accurate
# measurements.
def timed(fn):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    result = fn()
    end.record()
    torch.cuda.synchronize()
    return result, start.elapsed_time(end) / 1000


inp = torch.randn(4096, 4096).cuda()
print("compile:", timed(lambda: opt_foo3(inp))[1])
print("eager:", timed(lambda: foo3(inp))[1])

TRACED GRAPH
 ===== __compiled_fn_9_33a2ab81_c37c_4199_977d_b890f95c0d15 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[4096, 4096][4096, 1]cuda:0"):
        l_x_ = L_x_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:147 in foo3, code: y = x + 1
        y: "f32[4096, 4096][4096, 1]cuda:0" = l_x_ + 1;  l_x_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:148 in foo3, code: z = torch.nn.functional.relu(y)
        z: "f32[4096, 4096][4096, 1]cuda:0" = torch.nn.functional.relu(y);  y = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:149 in foo3, code: u = z * 2
        u: "f32[4096, 4096][4096, 1]cuda:0" = z * 2;  z = None
        return (u,)


compile: 0.5693040161132813
eager: 0.02673459243774414

请注意，与即时模式（eager mode）相比，torch.compile 完成任务似乎花费了更长时间。这是因为 torch.compile 在前几次执行时需要额外的时间来编译模型。torch.compile 会尽可能复用已编译的代码，因此如果我们多次运行优化后的模型，与即时模式相比，应该会看到显著的性能提升。

# turn off logging for now to prevent spam
torch._logging.set_logs(graph_code=False)

eager_times = []
for i in range(10):
    _, eager_time = timed(lambda: foo3(inp))
    eager_times.append(eager_time)
    print(f"eager time {i}: {eager_time}")
print("~" * 10)

compile_times = []
for i in range(10):
    _, compile_time = timed(lambda: opt_foo3(inp))
    compile_times.append(compile_time)
    print(f"compile time {i}: {compile_time}")
print("~" * 10)

import numpy as np

eager_med = np.median(eager_times)
compile_med = np.median(compile_times)
speedup = eager_med / compile_med
assert speedup > 1
print(
    f"(eval) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x"
)
print("~" * 10)

eager time 0: 0.0009103360176086426
eager time 1: 0.0008704000115394592
eager time 2: 0.0008785920143127441
eager time 3: 0.0008734719753265381
eager time 4: 0.0008704000115394592
eager time 5: 0.0008693760037422181
eager time 6: 0.0008724480271339416
eager time 7: 0.0008673279881477356
eager time 8: 0.0008704000115394592
eager time 9: 0.0008704000115394592
~~~~~~~~~~
compile time 0: 0.0005109760165214538
compile time 1: 0.0003624959886074066
compile time 2: 0.00036351999640464784
compile time 3: 0.00035836800932884214
compile time 4: 0.0003604480028152466
compile time 5: 0.00035942399501800537
compile time 6: 0.00036556801199913024
compile time 7: 0.00035737600922584536
compile time 8: 0.00035839998722076414
compile time 9: 0.00035942399501800537
~~~~~~~~~~
(eval) eager median: 0.0008704000115394592, compile median: 0.000359935998916626, speedup: 2.41820772070391x
~~~~~~~~~~

事实确实如此，我们可以看到使用 torch.compile 运行模型带来了显著的加速。加速主要源于减少了 Python 开销和 GPU 读写，因此观测到的加速效果可能会因模型架构和批大小等因素而异。例如，如果模型架构简单且数据量巨大，则瓶颈在于 GPU 计算，观测到的加速效果可能不那么明显。

要查看在真实模型上的加速效果，请查阅我们的 torch.compile 端到端教程。

相比 TorchScript 的优势 #

为什么我们应该使用 torch.compile 而不是 TorchScript？主要在于 torch.compile 能够以极少的代码改动处理任意 Python 代码。

与 TorchScript 相比，TorchScript 具有追踪模式（torch.jit.trace）和脚本模式（torch.jit.script）。追踪模式容易出现静默错误，而脚本模式需要重大的代码更改，并且在遇到不支持的 Python 代码时会引发错误。

例如，TorchScript 的追踪模式在处理依赖数据的控制流（下方的 if x.sum() < 0: 行）时会静默失败，因为它只追踪实际运行的控制流路径。相比之下，torch.compile 能够正确处理它。

def f1(x, y):
    if x.sum() < 0:
        return -y
    return y


# Test that `fn1` and `fn2` return the same result, given the same arguments `args`.
def test_fns(fn1, fn2, args):
    out1 = fn1(*args)
    out2 = fn2(*args)
    return torch.allclose(out1, out2)


inp1 = torch.randn(5, 5)
inp2 = torch.randn(5, 5)

traced_f1 = torch.jit.trace(f1, (inp1, inp2))
print("traced 1, 1:", test_fns(f1, traced_f1, (inp1, inp2)))
print("traced 1, 2:", test_fns(f1, traced_f1, (-inp1, inp2)))

compile_f1 = torch.compile(f1)
print("compile 1, 1:", test_fns(f1, compile_f1, (inp1, inp2)))
print("compile 1, 2:", test_fns(f1, compile_f1, (-inp1, inp2)))
print("~" * 10)

/var/lib/workspace/intermediate_source/torch_compile_tutorial.py:239: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if x.sum() < 0:
traced 1, 1: True
traced 1, 2: False
compile 1, 1: True
compile 1, 2: True
~~~~~~~~~~

TorchScript 的脚本模式可以处理依赖数据的控制流，但这通常需要重大的代码修改，并且在使用不支持的 Python 特性时会引发错误。

在下方的示例中，我们遗漏了 TorchScript 类型注解，从而收到了 TorchScript 错误，因为参数 y 的输入类型（int）与默认参数类型（torch.Tensor）不匹配。相比之下，torch.compile 在不需要任何类型注解的情况下即可工作。

import traceback as tb

torch._logging.set_logs(graph_code=True)


def f2(x, y):
    return x + y


inp1 = torch.randn(5, 5)
inp2 = 3

script_f2 = torch.jit.script(f2)
try:
    script_f2(inp1, inp2)
except:
    tb.print_exc()

compile_f2 = torch.compile(f2)
print("compile 2:", test_fns(f2, compile_f2, (inp1, inp2)))
print("~" * 10)

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_compile_tutorial.py", line 288, in <module>
    script_f2(inp1, inp2)
RuntimeError: f2() Expected a value of type 'Tensor (inferred)' for argument 'y' but instead found type 'int'.
Inferred 'y' to be of type 'Tensor' because it was not annotated with an explicit type.
Position: 1
Value: 3
Declaration: f2(Tensor x, Tensor y) -> Tensor
Cast error details: Unable to cast 3 to Tensor
TRACED GRAPH
 ===== __compiled_fn_18_5da28d0e_21da_42c8_a8ae_7e3b84baa970 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[5, 5][5, 1]cpu"):
        l_x_ = L_x_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:280 in f2, code: return x + y
        add: "f32[5, 5][5, 1]cpu" = l_x_ + 3;  l_x_ = None
        return (add,)


compile 2: True
~~~~~~~~~~

图中断是 torch.compile 中的核心概念之一。它允许 torch.compile 通过中断编译、运行不受支持的代码，然后恢复编译，从而处理任意的 Python 代码。“图中断”一词源于 torch.compile 试图捕获和优化 PyTorch 操作图的事实。当遇到不支持的 Python 代码时，该图必须被“中断”。图中断会导致优化机会的丢失，这虽然可能是不理想的，但比静默错误或直接崩溃要好。

让我们看一个依赖数据的控制流示例，以更好地了解图中断是如何工作的。

def bar(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b


opt_bar = torch.compile(bar)
inp1 = torch.ones(10)
inp2 = torch.ones(10)
opt_bar(inp1, inp2)
opt_bar(inp1, -inp2)

TRACED GRAPH
 ===== __compiled_fn_20_7e354d0e_91fd_430d_a418_3fb7201b06d9 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_a_: "f32[10][1]cpu", L_b_: "f32[10][1]cpu"):
        l_a_ = L_a_
        l_b_ = L_b_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:312 in bar, code: x = a / (torch.abs(a) + 1)
        abs_1: "f32[10][1]cpu" = torch.abs(l_a_)
        add: "f32[10][1]cpu" = abs_1 + 1;  abs_1 = None
        x: "f32[10][1]cpu" = l_a_ / add;  l_a_ = add = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:313 in bar, code: if b.sum() < 0:
        sum_1: "f32[][]cpu" = l_b_.sum();  l_b_ = None
        lt: "b8[][]cpu" = sum_1 < 0;  sum_1 = None
        return (lt, x)


TRACED GRAPH
 ===== __compiled_fn_24_7b5ce565_4267_453c_8db7_7c34516bd005 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[10][1]cpu", L_b_: "f32[10][1]cpu"):
        l_x_ = L_x_
        l_b_ = L_b_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:315 in torch_dynamo_resume_in_bar_at_313, code: return x * b
        mul: "f32[10][1]cpu" = l_x_ * l_b_;  l_x_ = l_b_ = None
        return (mul,)


TRACED GRAPH
 ===== __compiled_fn_26_c5112331_f131_4389_a230_825c45f8871f =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_b_: "f32[10][1]cpu", L_x_: "f32[10][1]cpu"):
        l_b_ = L_b_
        l_x_ = L_x_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:314 in torch_dynamo_resume_in_bar_at_313, code: b = b * -1
        b: "f32[10][1]cpu" = l_b_ * -1;  l_b_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:315 in torch_dynamo_resume_in_bar_at_313, code: return x * b
        mul_1: "f32[10][1]cpu" = l_x_ * b;  l_x_ = b = None
        return (mul_1,)



tensor([0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
        0.5000])

我们第一次运行 bar 时，可以看到 torch.compile 追踪到了 2 个图，对应于以下代码（注意此时 b.sum() < 0 为 False）：

x = a / (torch.abs(a) + 1); b.sum()
return x * b

第二次运行 bar 时，我们走入 if 语句的另一个分支，得到 1 个被追踪的图，对应于代码 b = b * -1; return x * b。我们没有看到 x = a / (torch.abs(a) + 1); b.sum() 的图再次输出，因为 torch.compile 从第一次运行中缓存并复用了该图。

让我们通过示例研究 TorchDynamo 是如何遍历 bar 的。如果 b.sum() < 0，那么 TorchDynamo 将运行图 1，让 Python 确定条件判断的结果，然后运行图 2。另一方面，如果 not b.sum() < 0，那么 TorchDynamo 将运行图 1，让 Python 确定条件判断的结果，然后运行图 3。

我们可以通过 torch._logging.set_logs(graph_breaks=True) 查看所有的图中断。

# Reset to clear the torch.compile cache
torch._dynamo.reset()
opt_bar(inp1, inp2)
opt_bar(inp1, -inp2)

TRACED GRAPH
 ===== __compiled_fn_28_a2431937_6d03_42a4_baeb_4ad752cab0a4 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_a_: "f32[10][1]cpu", L_b_: "f32[10][1]cpu"):
        l_a_ = L_a_
        l_b_ = L_b_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:312 in bar, code: x = a / (torch.abs(a) + 1)
        abs_1: "f32[10][1]cpu" = torch.abs(l_a_)
        add: "f32[10][1]cpu" = abs_1 + 1;  abs_1 = None
        x: "f32[10][1]cpu" = l_a_ / add;  l_a_ = add = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:313 in bar, code: if b.sum() < 0:
        sum_1: "f32[][]cpu" = l_b_.sum();  l_b_ = None
        lt: "b8[][]cpu" = sum_1 < 0;  sum_1 = None
        return (lt, x)


TRACED GRAPH
 ===== __compiled_fn_32_e5ddb3c6_62b8_4ed7_af4c_64dcb34e9998 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_x_: "f32[10][1]cpu", L_b_: "f32[10][1]cpu"):
        l_x_ = L_x_
        l_b_ = L_b_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:315 in torch_dynamo_resume_in_bar_at_313, code: return x * b
        mul: "f32[10][1]cpu" = l_x_ * l_b_;  l_x_ = l_b_ = None
        return (mul,)


TRACED GRAPH
 ===== __compiled_fn_34_83e25ea1_e672_4f1d_af87_299ddec808d5 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_b_: "f32[10][1]cpu", L_x_: "f32[10][1]cpu"):
        l_b_ = L_b_
        l_x_ = L_x_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:314 in torch_dynamo_resume_in_bar_at_313, code: b = b * -1
        b: "f32[10][1]cpu" = l_b_ * -1;  l_b_ = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:315 in torch_dynamo_resume_in_bar_at_313, code: return x * b
        mul_1: "f32[10][1]cpu" = l_x_ * b;  l_x_ = b = None
        return (mul_1,)



tensor([0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
        0.5000])

为了最大化加速效果，应当限制图中断。我们可以通过使用 fullgraph=True，强制 TorchDynamo 在遇到第一次图中断时就引发错误。

# Reset to clear the torch.compile cache
torch._dynamo.reset()

opt_bar_fullgraph = torch.compile(bar, fullgraph=True)
try:
    opt_bar_fullgraph(torch.randn(10), torch.randn(10))
except:
    tb.print_exc()

Traceback (most recent call last):
  File "/var/lib/workspace/intermediate_source/torch_compile_tutorial.py", line 360, in <module>
    opt_bar_fullgraph(torch.randn(10), torch.randn(10))
  File "/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 1177, in compile_wrapper
    raise e.with_traceback(
torch._dynamo.exc.Unsupported: Data-dependent branching
  Explanation: Detected data-dependent branching (e.g. `if my_tensor.sum() > 0:`). Dynamo does not support tracing dynamic control flow.

      The branch condition involves a tensor computed as follows:
        # File "/var/lib/workspace/intermediate_source/torch_compile_tutorial.py", line 313, in bar, code: if b.sum() < 0:
        lt = lt(sum_1, 0)

  Hint: For the common pattern `if tensor_cond: x = transform(x)` (e.g. clamping inf/nan values), consider making the code branchless by always applying the transform. Operations like torch.clamp, torch.nan_to_num, and torch.where are typically no-ops on well-behaved inputs and compile without graph breaks.
  Hint: This graph break is fundamental - it is unlikely that Dynamo will ever be able to trace through your code. Consider finding a workaround.
  Hint: Use `torch.cond` to express dynamic control flow.

  Developer debug context: attempted to jump with TensorVariable()

 For more details about this graph break, please visit: https://meta-pytorch.cn/compile-graph-break-site/gb/gb0170.html

from user code:
   File "/var/lib/workspace/intermediate_source/torch_compile_tutorial.py", line 313, in bar
    if b.sum() < 0:

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

在上面的例子中，我们可以通过用 torch.cond 替换 if 语句来规避这种图中断。

from functorch.experimental.control_flow import cond


@torch.compile(fullgraph=True)
def bar_fixed(a, b):
    x = a / (torch.abs(a) + 1)

    def true_branch(y):
        return y * -1

    def false_branch(y):
        # NOTE: torch.cond doesn't allow aliased outputs
        return y.clone()

    b = cond(b.sum() < 0, true_branch, false_branch, (b,))
    return x * b


bar_fixed(inp1, inp2)
bar_fixed(inp1, -inp2)

TRACED GRAPH
 ===== __compiled_fn_37_f536baee_f0cd_41b0_a71e_d72d2403f910 =====
 /var/lib/ci-user/.local/lib/python3.10/site-packages/torch/fx/_lazy_graph_module.py class GraphModule(torch.nn.Module):
    def forward(self, L_a_: "f32[10][1]cpu", L_b_: "f32[10][1]cpu"):
        l_a_ = L_a_
        l_b_ = L_b_

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:373 in bar_fixed, code: x = a / (torch.abs(a) + 1)
        abs_1: "f32[10][1]cpu" = torch.abs(l_a_)
        add: "f32[10][1]cpu" = abs_1 + 1;  abs_1 = None
        x: "f32[10][1]cpu" = l_a_ / add;  l_a_ = add = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:382 in bar_fixed, code: b = cond(b.sum() < 0, true_branch, false_branch, (b,))
        sum_1: "f32[][]cpu" = l_b_.sum()
        lt: "b8[][]cpu" = sum_1 < 0;  sum_1 = None
        cond_true_0 = self.cond_true_0
        cond_false_0 = self.cond_false_0
        cond = torch.ops.higher_order.cond(lt, cond_true_0, cond_false_0, (l_b_,));  lt = cond_true_0 = cond_false_0 = l_b_ = None
        b: "f32[10][1]cpu" = cond[0];  cond = None

        # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:383 in bar_fixed, code: return x * b
        mul: "f32[10][1]cpu" = x * b;  x = b = None
        return (mul,)

    class cond_true_0(torch.nn.Module):
        def forward(self, l_b_: "f32[10][1]cpu"):
            l_b__1 = l_b_

            # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:376 in true_branch, code: return y * -1
            mul: "f32[10][1]cpu" = l_b__1 * -1;  l_b__1 = None
            return (mul,)

    class cond_false_0(torch.nn.Module):
        def forward(self, l_b_: "f32[10][1]cpu"):
            l_b__1 = l_b_

            # File: /var/lib/workspace/intermediate_source/torch_compile_tutorial.py:380 in false_branch, code: return y.clone()
            clone: "f32[10][1]cpu" = l_b__1.clone();  l_b__1 = None
            return (clone,)



tensor([0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
        0.5000])

如果需要序列化图或在不同环境（即无 Python 环境）中运行图，请考虑改用 torch.export（自 PyTorch 2.1+ 起）。一个重要的限制是 torch.export 不支持图中断。请查阅 torch.export 教程以了解关于 torch.export 的更多详细信息。

查阅 torch.compile 编程模型中关于图中断的章节，获取关于如何规避图中断的建议。

故障排查 #

torch.compile 未能加速您的模型？编译时间过长？代码在过度重新编译？处理图中断时遇到困难？正在寻找关于如何最好地使用 torch.compile 的技巧？或者只是想了解更多关于 torch.compile 的内部运作原理？

请查阅 torch.compile 编程模型。

结论 #

在本教程中，我们介绍了 torch.compile，涵盖了基本用法、演示了对比即时模式的加速效果、与 TorchScript 进行了对比，并简要描述了图中断。

如需查看关于真实模型的端到端示例，请查阅我们的 torch.compile 端到端教程。

如需排查问题并深入了解如何将 torch.compile 应用于您的代码，请查阅 torch.compile 编程模型。

希望您能尝试一下 torch.compile！

脚本总运行时间：（0 分 14.489 秒）

`torch.compile` 简介#

基本用法 #

演示加速效果 #

相比 TorchScript 的优势 #

图中断（Graph Breaks）#

故障排查 #

结论 #

文档

教程

资源

torch.compile 简介#

基本用法#

演示加速效果#

相比 TorchScript 的优势#

图中断（Graph Breaks）#

故障排查#

结论#

文档

教程

资源

`torch.compile` 简介#

基本用法 #

演示加速效果 #

相比 TorchScript 的优势 #

故障排查 #

结论 #