(beta) 使用火炬日志 Python API 与 torch.compile 指令

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：迈克尔·拉佐斯

import logging

本教程介绍了 TORCH_LOGS 环境变量，以及 Python API，并演示了如何将其应用于观察 torch.compile 的阶段。

备注

本教程需要 PyTorch 2.2.0 或更高版本。

设置

在本例中，我们将设置一个简单的 Python 函数，该函数执行逐元素加法，并使用 TORCH_LOGS Python API 观察编译过程。

备注

也有一个环境变量 TORCH_LOGS ，可以用于在命令行中更改日志设置。每个示例都显示了等效的环境变量设置。

import torch

# exit cleanly if we are on a device that doesn't support torch.compile
if torch.cuda.get_device_capability() < (7, 0):
    print("Skipping because torch.compile is not supported on this device.")
else:
    @torch.compile()
    def fn(x, y):
        z = x + y
        return z + 2


    inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))


# print separator and reset dynamo
# between each example
    def separator(name):
        print(f"==================={name}=========================")
        torch._dynamo.reset()


    separator("Dynamo Tracing")
# View dynamo tracing
# TORCH_LOGS="+dynamo"
    torch._logging.set_logs(dynamo=logging.DEBUG)
    fn(*inputs)

    separator("Traced Graph")
# View traced graph
# TORCH_LOGS="graph"
    torch._logging.set_logs(graph=True)
    fn(*inputs)

    separator("Fusion Decisions")
# View fusion decisions
# TORCH_LOGS="fusion"
    torch._logging.set_logs(fusion=True)
    fn(*inputs)

    separator("Output Code")
# View output code generated by inductor
# TORCH_LOGS="output_code"
    torch._logging.set_logs(output_code=True)
    fn(*inputs)

    separator("")

结论 ¶

在本教程中，我们通过实验少量可用的日志选项，介绍了 TORCH_LOGS 环境变量和 Python API。要查看所有可用选项的描述，请运行任何导入 torch 的 Python 脚本并将 TORCH_LOGS 设置为“help”。

或者，您也可以查看 torch._logging 文档以查看所有可用日志选项的描述。

有关 torch.compile 的更多信息，请参阅 torch.compile 教程。