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

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 上 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 性能提升约 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 是一个值得投入的工具。