PyTorch JIT 与 TorchScript:实测推理性能提升 50%
PyTorch 主要支持两种运行模式:Eager Mode(急切模式)和 Script Mode(脚本模式)。Eager 模式适合模型编写、训练和调试,而 Script 模式专为部署设计。Script 模式的核心在于 PyTorch JIT 编译器和 TorchScript 格式。
核心概念
Script 模式通过 torch.jit.trace 或 torch.jit.script 将 PyTorch eager module 转换为中间表示(IR)。这个 IR 经过内部优化,并在运行时由 PyTorch JIT 编译器处理。关键在于,该 IR 与 Python 运行时解耦,这意味着模型可以在没有 Python 解释器的环境中运行。
为什么要用 Script 模式?
- 脱离 Python 限制:可以绕过 Python GIL 和运行时开销,例如通过 LibTorch 在 C++ 中运行模型,非常适合 IoT 等嵌入式场景。
- 性能优化:JIT 编译器利用运行时信息进行优化,如层融合(Layer Fusion)、量化(Quantization)和稀疏化。通常比原生 PyTorch 模型更快。
- 跨平台部署:支持导出为独立文件,甚至可转换为 ONNX 格式在其他框架中运行。
- 资源受限环境:针对内存和计算资源受限的设备进行优化。
Trace 与 Script 的区别
这两个函数都是将 Python 代码转换为 TorchScript 的方法,但适用场景不同:
torch.jit.trace:基于特定输入(通常是张量)跟踪执行过程,记录计算图。适用于结构固定、输入大小不变的模型(如预训练模型)。torch.jit.script:直接解析 Python 语法并编译。适用于动态图模型,结构和输入可在运行时变化(如 RNN 或变长序列)。
注意:虽然两者各有优势,但在静态图场景下,
trace往往更便捷;而在需要动态控制流时,script是必须的。
性能实测
为了验证性能差异,我们分别在 CPU 和 GPU 上对比了 Eager 模式和 TorchScript 模式的推理速度。测试硬件为 Google Colab(Intel Xeon @ 2.00GHz, Tesla T4)。
1. BERT 模型测试
使用 HuggingFace 的 BERT 模型作为示例。
from transformers import BertTokenizer, BertModel
import numpy as np
import torch
from time import perf_counter
def timer(f, *args):
start = perf_counter()
f(*args)
return (1000 * (perf_counter() - start))
# 加载原生模型
native_model = BertModel.from_pretrained("bert-base-uncased")
script_model = BertModel.from_pretrained(, torchscript=)
script_tokenizer = BertTokenizer.from_pretrained(, torchscript=)
text =
tokenized_text = script_tokenizer.tokenize(text)
masked_index =
tokenized_text[masked_index] =
indexed_tokens = script_tokenizer.convert_tokens_to_ids(tokenized_text)
segments_ids = [, , , , , , , , , , , , , ]
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])

