1. PyTorch模型导出基础：理解TorchScript、JIT trace与script的区别，导出第一个TorchScript模型

模型部署的第一步，就是把训练好的PyTorch模型「搬」到生产环境。但这里有个坑——PyTorch模型本身依赖Python运行时，而生产环境往往没有Python解释器，或者只有C++运行时。怎么办？

TorchScript就是PyTorch给出的答案。它本质上是一种序列化格式，能把你的模型从Python中「解放」出来。我刚开始接触部署时，总觉得这步多此一举，直到有一次线上服务因为Python版本冲突挂了……嗯，从那以后我再也不敢跳过TorchScript了。

1.1 什么是TorchScript？

TorchScript是PyTorch模型的一种中间表示。它不依赖Python解释器，可以被C++运行时加载执行。说白了，就是给模型拍了一张「快照」——把模型的结构和参数都固定下来。

它的核心价值有三点：

  • 可移植性：脱离Python环境运行，C++、移动端、嵌入式都能用
  • 性能优化：支持JIT编译，推理速度比原生PyTorch快不少
  • 序列化：保存为.pt文件，方便分发和版本管理

重要概念：TorchScript不是一个新的框架，而是PyTorch模型的一种「固化」形式。你训练好的模型，通过TorchScript导出后，就不再需要Python了。

1.2 JIT trace vs script：两种导出方式

PyTorch提供了两种导出TorchScript的方式：torch.jit.trace 和 torch.jit.script。它们各有适用场景，我建议你根据模型结构来选择。

特性 torch.jit.trace torch.jit.script
原理 跟踪实际执行路径 解析Python源码
适用场景 静态图、无控制流 动态图、有控制流
优点 简单、速度快 支持if/for等控制流
缺点 不支持动态分支 语法限制较多
调试难度 低 中等

1.3 trace：跟踪执行路径

torch.jit.trace 的工作方式很简单：你给模型一个示例输入，它「跟踪」模型在这个输入下的执行路径，然后把所有操作记录下来。

为什么会这样？因为trace只记录实际执行过的操作。如果你的模型里有 if x > 0 这样的分支，trace只会记录当前输入走的那条路。这意味着——

注意：trace不支持动态控制流。如果你的模型里有条件分支、循环等，trace可能会导出错误的图。我曾经在项目中用trace导出一个带if-else的模型，结果线上推理时永远只走一个分支……排查了半天才发现是trace的问题。

来看一个trace的示例：

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)
    
    def forward(self, x):
        return self.fc(x)

model = SimpleModel()
example_input = torch.randn(1, 10)

# 使用trace导出
traced_model = torch.jit.trace(model, example_input)
traced_model.save("simple_model.pt")

# 加载并推理
loaded_model = torch.jit.load("simple_model.pt")
output = loaded_model(example_input)
print(output)

这段代码很简单。但注意，trace要求你提供正确的示例输入——输入的形状、数据类型必须和实际推理时一致。否则导出的模型可能无法正常工作。

1.4 script：解析Python源码

如果你的模型里有控制流，那就得用 torch.jit.script 了。它会解析你的Python源码，把整个模型「翻译」成TorchScript。

script的优点是支持动态控制流，但代价是语法限制较多。比如，你不能用Python的某些动态特性（如动态类型、某些内置函数）。

我个人习惯是：能用trace就用trace，实在不行再上script。因为script的调试成本更高，有时候一个简单的语法问题就能卡你半天。

import torch
import torch.nn as nn

class DynamicModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(10, 5)
    
    def forward(self, x, use_fc1: bool):
        if use_fc1:
            return self.fc1(x)
        else:
            return self.fc2(x)

model = DynamicModel()

# 使用script导出
scripted_model = torch.jit.script(model)
scripted_model.save("dynamic_model.pt")

# 加载并推理
loaded_model = torch.jit.load("dynamic_model.pt")
output = loaded_model(torch.randn(1, 10), True)
print(output)

小技巧：使用script时，建议给函数的参数加上类型注解。比如上面的 use_fc1: bool，这能帮助TorchScript正确解析你的代码。我在项目中遇到过因为缺少类型注解导致导出失败的情况，加上就好了。

1.5 trace vs script：如何选择？

你想想看，实际项目中大部分模型都是静态的——没有控制流，输入输出形状固定。这种情况下，trace完全够用。

但如果你遇到以下情况，就得用script了：

  • 模型里有 if、for、while 等控制流
  • 模型依赖输入的形状或值来决定执行路径
  • 模型使用了某些Python特性（如字典、列表操作）

还有一个折中方案：先用trace导出，如果报错再改用script。我经常这么干，因为trace的报错信息通常更清晰，能帮你快速定位问题。

1.6 导出第一个TorchScript模型

好了，理论说完了，我们来动手导出第一个TorchScript模型。我建议你从最简单的线性模型开始，这样能快速理解整个流程。

import torch
import torch.nn as nn
import torch.nn.functional as F

class FirstModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.fc = nn.Linear(64 * 6 * 6, 10)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = FirstModel()
model.eval()  # 记得切换到评估模式

# 创建示例输入
example_input = torch.randn(1, 1, 28, 28)

# 导出
traced_model = torch.jit.trace(model, example_input)
traced_model.save("first_model.pt")

# 验证
loaded_model = torch.jit.load("first_model.pt")
with torch.no_grad():
    output = loaded_model(example_input)
    print(f"输出形状: {output.shape}")
    print(f"输出值: {output}")

关键点：导出前一定要调用 model.eval()。否则模型会处于训练模式，Dropout、BatchNorm等层的行为会不同，导致导出的模型推理结果不正确。我刚开始做部署时就踩过这个坑，导出的模型在C++端推理结果和Python端对不上，排查了好久才发现是忘了切模式。

1.7 常见问题与避坑指南

导出TorchScript时，有几个常见问题需要注意：

  • 输入形状不匹配：trace时提供的示例输入形状，必须和实际推理时一致。如果输入形状会变化，建议用script或者固定输入尺寸。
  • 不支持的操作：某些PyTorch操作在TorchScript中不支持。比如 torch.Tensor.numpy()、torch.Tensor.item() 等。遇到这种情况，可以用 torch.jit.is_scripting() 来做条件判断。
  • 动态控制流：如果模型里有动态控制流，trace会导出错误的图。用script可以解决，但要注意语法限制。

避坑指南：我曾经在项目中用trace导出一个带 torch.argmax 的模型，结果导出后推理结果全错了。原因是 torch.argmax 在trace时被「固化」了——它只记录了第一次推理时的索引值。后来改用script才解决。所以，如果你的模型里有类似的操作，建议优先用script。

1.8 小结

这一章我们聊了TorchScript的基础知识，包括：

  • TorchScript是什么，为什么需要它
  • trace和script的区别与选择
  • 如何导出第一个TorchScript模型
  • 常见问题与避坑指南

下一章，我们会深入探讨TorchScript的高级用法，包括如何调试导出的模型、如何处理动态输入、以及如何优化推理性能。嗯，到时候见。

← 上一章 📖 返回目录 下一章 →