PyTorch JIT 与 TorchScript：实测推理性能提升 50%

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 模式？

1. 脱离 Python 限制：可以绕过 Python GIL 和运行时开销，例如通过 LibTorch 在 C++ 中运行模型，非常适合 IoT 等嵌入式场景。
2. 性能优化：JIT 编译器利用运行时信息进行优化，如层融合（Layer Fusion）、量化（Quantization）和稀疏化。通常比原生 PyTorch 模型更快。
3. 跨平台部署：支持导出为独立文件，甚至可转换为 ONNX 格式在其他框架中运行。
4. 资源受限环境：针对内存和计算资源受限的设备进行优化。

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")

# 加载 TorchScript 模型
script_model = BertModel.from_pretrained("bert-base-uncased", torchscript=True)
script_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', torchscript=True)

# 准备输入数据
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = script_tokenizer.tokenize(text)
masked_index = 8
tokenized_text[masked_index] = '[MASK]'
indexed_tokens = script_tokenizer.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])

CPU 测试结果：

native_model.eval()
print(np.mean([timer(native_model, tokens_tensor, segments_tensors) for _ in range(100)]))

GPU 测试结果：

native_model = native_model.cuda()
tokens_tensor_gpu = tokens_tensor.cuda()
segments_tensors_gpu = segments_tensors.cuda()
native_model.eval()
print(np.mean([timer(native_model, tokens_tensor_gpu, segments_tensors_gpu) for _ in range(100)]))

对 TorchScript 模型同样进行测试（注意：trace 后无需显式调用 .eval()）：

traced_model = torch.jit.trace(script_model, [tokens_tensor, segments_tensors])
print(np.mean([timer(traced_model, tokens_tensor, segments_tensors) for _ in range(100)]))

结果对比：

模式	CPU Latency (ms)	GPU Latency (ms)
PyTorch (Eager)	171.27	30.42
TorchScript	165.24	13.50

数据显示，在 CPU 上 TorchScript 提升了约 3.5%，而在 GPU 上提升了 55.6%。

2. ResNet 模型测试

为了验证结论的普适性，我们使用了 ResNet18 再次测试。

import torchvision
import torch
from time import perf_counter
import numpy as np

def timer(f, *args):
    start = perf_counter()
    f(*args)
    return (1000 * (perf_counter() - start))

# PyTorch CPU 版本
model_ft = torchvision.models.resnet18(pretrained=True)
model_ft.eval()
x_ft = torch.rand(1, 3, 224, 224)
print(f'pytorch cpu: {np.mean([timer(model_ft, x_ft) for _ in range(10)])}')

# PyTorch GPU 版本
model_ft_gpu = torchvision.models.resnet18(pretrained=True).cuda()
x_ft_gpu = x_ft.cuda()
model_ft_gpu.eval()
print(f'pytorch gpu: {np.mean([timer(model_ft_gpu, x_ft_gpu) for _ in range(10)])}')

# TorchScript CPU 版本
script_cell = torch.jit.script(model_ft, (x_ft))
print(f'torchscript cpu: {np.mean([timer(script_cell, x_ft) for _ in range(10)])}')

# TorchScript GPU 版本
script_cell_gpu = torch.jit.script(model_ft_gpu, (x_ft_gpu))
print(f'torchscript gpu: {np.mean([timer(script_cell_gpu, x_ft.cuda()) for _ in range(100)])}')

结果对比：

模式	CPU Latency (ms)	GPU Latency (ms)
PyTorch (Eager)	77.47	2.99
TorchScript	74.24	1.64

结论一致：CPU 性能提升约 4.2%，GPU 性能提升约 45%。

总结

本次测试重点对比了 PyTorch 的 Eager 模式和 Script 模式。通过 BERT 和 ResNet 两个典型网络在 CPU 和 GPU 上的表现，我们发现：

1. 部署效率更高：TorchScript 允许模型脱离 Python 环境运行，便于集成到 C++ 服务或嵌入式设备中。
2. 推理性能显著提升：在 GPU 环境下，TorchScript 相比原生 PyTorch 能获得 50% 左右的推理加速；CPU 环境下也有 4% 左右的提升。
3. 选择建议：对于静态结构模型，trace 足够高效；涉及动态控制流时，请使用 script。

如果你需要在生产环境中追求极致的推理性能，TorchScript 是一个值得投入的工具。