TorchScript 编译器详解：如何通过 trace 与 script 将动态图导出为高性能静态图

TorchScript 是 PyTorch 官方提供的编译器，用于将 PyTorch 模型从灵活的 Python 动态图环境转换成高性能的静态图表示。这种静态图格式可以脱离 Python 解释器运行，实现推理加速，并支持在 C++ 或其他生产环境中部署，是模型部署的关键技术。

本文将深入解析 TorchScript 的两种核心转换方法：追踪（Tracing） 和 脚本化（Scripting），并提供实操代码。

为什么需要 TorchScript？

PyTorch 默认使用动态图，这意味着计算图是实时构建和执行的。这给开发带来了极大的灵活性，但在生产部署时，动态图可能导致额外的开销。通过 TorchScript 转换为静态图后，可以实现以下优势：

1. 推理加速： 消除 Python 解释器带来的开销。
2. C++ 部署： 使用 LibTorch 库在 C++ 应用程序中加载和运行模型。
3. 跨平台/语言支持： 方便模型在各种环境中进行序列化和部署。

方法一：追踪 (Tracing) – torch.jit.trace

追踪是最简单、最常用的转换方法。它通过提供一组示例输入数据，记录模型在执行这些输入时的确切操作序列，从而构建一个静态图。

工作原理

当 torch.jit.trace 运行时，它实际上执行了一次 forward 函数，并记录下所有被调用的操作（Ops）和数据流。它记录的是执行路径，而不是代码本身。

限制

由于它只记录执行路径，追踪无法正确处理依赖于输入数据或内部状态的控制流（如 if/else 语句、动态循环）。它只会记录执行时实际走的那个分支。

实践代码：简单的模型追踪

我们创建一个简单的线性层模型并对其进行追踪：

import torch

# 1. 定义一个简单的模型
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = torch.nn.Linear(10, 5)

    def forward(self, x):
        # 追踪适合这种顺序执行的结构
        return torch.relu(self.fc(x))

# 实例化模型并设置为评估模式
model = SimpleModel().eval()
# 准备示例输入 (Batch size 1, Feature size 10)
example_input = torch.randn(1, 10)

# 2. 使用 torch.jit.trace 追踪模型
traced_model = torch.jit.trace(model, example_input)

print("--- 追踪模型的图结构 (简化版) ---")
# 打印追踪到的静态计算图
print(traced_model.graph)

# 3. 保存追踪后的模型
# traced_model.save("simple_traced_model.pt")
# print("模型已保存为 simple_traced_model.pt")

# 4. 验证结果一致性
original_output = model(example_input)
traced_output = traced_model(example_input)
print(f"输出差异是否极小: {torch.allclose(original_output, traced_output)}")

方法二：脚本化 (Scripting) – torch.jit.script

脚本化是通过直接解析 Python 源代码，将其编译成 TorchScript 语言的 IR（中间表示）来实现的。这是处理复杂模型和控制流的推荐方法。

工作原理

torch.jit.script 像传统的编译器一样工作，它会分析函数的完整结构，包括所有的 if/else 块、循环和变量赋值。它不依赖于实际的输入值。

限制

脚本化要求模型代码遵循 TorchScript 兼容的 Python 子集。例如，不能使用 Python 列表进行复杂的动态操作，也不能使用某些依赖于 Python 解释器的高级特性。

实践代码：处理控制流

我们创建一个包含条件判断的模型，并使用脚本化进行转换：

import torch

# 1. 定义一个包含控制流的模型
class DynamicModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 定义一个参数，其值决定走哪个分支
        self.flag = torch.nn.Parameter(torch.tensor(0.5))

    # 使用 @torch.jit.script 方法修饰类或者函数
    # 也可以直接调用 torch.jit.script(Model()) 来转换整个实例
    def forward(self, x):
        # 脚本化可以正确捕获 if/else 结构
        if self.flag > 0.0:
            print("--- 路径 A: 乘 2 ---")
            return x * 2.0
        else:
            print("--- 路径 B: 除 2 ---")
            return x / 2.0

model_dynamic = DynamicModel().eval()
example_input = torch.randn(1, 3)

# 2. 使用 torch.jit.script 脚本化
# 注意：此处不会执行 forward，而是直接解析代码
scripted_model = torch.jit.script(model_dynamic)

# 3. 验证脚本化后的模型
print("\n--- 执行脚本化模型 ---")
# 脚本化模型执行时，可以正确地保留 if 逻辑
output_script = scripted_model(example_input)

# 4. 修改模型参数并再次执行 (验证静态图包含两种逻辑)
model_dynamic.flag.data = torch.tensor(-1.0)
scripted_output_new = scripted_model(example_input)

# scripted_model.save("dynamic_scripted_model.pt")
# print("模型已保存为 dynamic_scripted_model.pt")

print(f"Flag > 0.0 时的输出 (路径 A): {output_script.mean():.4f}")
print(f"Flag < 0.0 时的输出 (路径 B): {scripted_output_new.mean():.4f}")

总结与选择

建议：

1. 对于像 ResNet 或 VGG 这种纯粹的层堆叠模型，优先使用 Tracing。
2. 如果模型中包含复杂的逻辑判断、列表操作或自定义的控制流，必须使用 Scripting。