PyTorch支持两种模式：eager模式和script模式。eager模式主要用于模型的编写、训练和调试，script模式主要是针对部署的，其包含PytorchJIT和TorchScript（一种在 PyTorch 中执行高效的序列化代码格式）。
script模式使用torch.jit.trace和torch.jit.script创建一个PyTorch eager module的中间表示（intermediate representation, IR），IR 经过内部优化，并在运行时使用 PyTorch JIT 编译。PyTorch JIT 编译器使用运行时信息来优化 IR。该 IR 与 Python 运行时是解耦的。
PyTorch JIT（Just-In-Time Compilation）是 PyTorch 中的即时编译器。
- 它允许你将模型转化为 TorchScript 格式，从而提高模型的性能和部署效率。
- JIT 允许你在动态图和静态图之间无缝切换。你可以在 Python 中以动态图的方式构建和调试模型，然后将模型编译为 TorchScript 以进行优化和部署。
- JIT 允许你在不同的深度学习框架之间进行模型转换，例如将 PyTorch 模型转换为 ONNX 格式，从而可以在其他框架中运行。
TorchScript 是 PyTorch 提供的一种将模型序列化以便在其他环境中运行的机制。它将 PyTorch 模型编译成一种中间表示形式，可以在没有 Python 解释器的环境中运行。这使得模型可以在 C++ 等其他语言中运行，也可以在嵌入式设备等资源受限的环境中实现高效的推理。
以下是 TorchScript 的一些重要特性和用途：
- 静态图表示形式：TorchScript 是一种静态图表示形式，它在模型构建阶段对计算图进行编译和优化，而不是在运行时动态构建。这可以提高模型的执行效率。
- 模型导出：TorchScript 允许将 PyTorch 模型导出到一个独立的文件中，然后可以在没有 Python 环境的设备上运行。
- 跨平台部署：TorchScript 允许在不同的深度学习框架之间进行模型转换，例如将 PyTorch 模型转换为 ONNX 格式，从而可以在其他框架中运行。
- 模型优化和量化：通过 TorchScript，你可以使用各种技术（如量化）对模型进行优化，从而减小模型的内存占用和计算资源消耗。
- 融合和集成：TorchScript 可以帮助你将多个模型整合到一个整体流程中，从而提高系统的整体性能。
- 嵌入式设备：对于资源受限的嵌入式设备，TorchScript 可以帮助你优化模型以适应这些环境。
使用 TorchScript 可以将 PyTorch 模型变得更容易在生产环境中部署和集成。然而，它也可能需要你对模型进行一些修改以使其可以成功编译为 TorchScript。
总的来说，TorchScript 是一个强大的工具，特别是对于需要在不同环境中部署 PyTorch 模型的情况。通过将模型导出为 TorchScript，你可以实现更广泛的模型应用和部署。
一段话总结，为什么要用以及什么时候要用script模式呢？
- 可以脱离python GIL以及python runtime的限制来运行模型，比如通过LibTorch通过C++来运行模型。这样方便了模型部署，例如可以在IoT等平台上运行。例如这个tutorial，使用C++来运行pytorch的model。
- PyTorch JIT是用于pytorch的优化的JIT编译器，它使用运行时信息来优化 TorchScript modules，可以自动进行层融合、量化、稀疏化等优化。因此，相比pytorch model，TorchScript的性能会更高。
Script mode通过torch.jit.trace或者torch.jit.script来调用。这两个函数都是将python代码转换为TorchScript的两种不同的方法。torch.jit.trace将一个特定的输入（通常是一个张量，需要我们提供一个input）传递给一个PyTorch模型，torch.jit.trace会跟踪此input在model中的计算过程，然后将其转换为Torch脚本。这个方法适用于那些在静态图中可以完全定义的模型，例如具有固定输入大小的神经网络。通常用于转换预训练模型。torch.jit.script直接将Python函数（或者一个Python模块）通过python语法规则和编译转换为Torch脚本。torch.jit.script更适用于动态图模型，这些模型的结构和输入可以在运行时发生变化。例如，对于RNN或者一些具有可变序列长度的模型，使用torch.jit.script会更为方便。
在通常情况下，更应该倾向于使用torch.jit.trace而不是torch.jit.script。
在上一篇blog中，我们非常非常详细介绍了torch.jit.trace和torch.jit.script的区别以及使用建议。强烈建议先阅读上一篇blog，再来阅读此篇内容。
本篇中，我们重点看一下TorchScript model与eager model的性能区别。
JIT Trace
torch.jit.trace使用eager model和一个dummy input作为输入，tracer会根据提供的model和input记录数据在模型中的流动过程，然后将整个模型转换为TorchScript module。看一个具体的例子：
我们使用BERT（Bidirectional Encoder Representations from Transformers）作为例子。
from transformers import BertTokenizer, BertModel
import numpy as np
import torch
from time import perf_counterdef timer(f,*args): start = perf_counter()f(*args)return (1000 * (perf_counter() - start))# 加载bert model
native_model = BertModel.from_pretrained("bert-base-uncased")
# huggingface的API中，使用torchscript=True参数可以直接加载TorchScript model
script_model = BertModel.from_pretrained("bert-base-uncased", torchscript=True)script_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', torchscript=True)# Tokenizing input text
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = script_tokenizer.tokenize(text)# Masking one of the input tokens
masked_index = 8tokenized_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]# Creating a dummy input
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])
然后分别在CPU和GPU上测试eager mode的pytorch推理速度。
# 在CPU上测试eager model推理性能
native_model.eval()
np.mean([timer(native_model,tokens_tensor,segments_tensors) for _ in range(100)])# 在GPU上测试eager model推理性能
native_model = native_model.cuda()
native_model.eval()
tokens_tensor_gpu = tokens_tensor.cuda()
segments_tensors_gpu = segments_tensors.cuda()
np.mean([timer(native_model,tokens_tensor_gpu,segments_tensors_gpu) for _ in range(100)])
再分别在CPU和GPU上测试script mode的TorchScript模型的推理速度
# 在CPU上测试TorchScript性能
traced_model = torch.jit.trace(script_model, [tokens_tensor, segments_tensors])
# 因模型的trace时，已经包含了.eval()的行为，因此不必再去显式调用model.eval()
np.mean([timer(traced_model,tokens_tensor,segments_tensors) for _ in range(100)])# 在GPU上测试TorchScript的性能最终运行结果如表
| CPU latency (ms) | GPU latency (ms) | |
|---|---|---|
| PyTorch | 171.27 | 30.42 |
| TorchScript | 165.24 | 13.50 |
我使用的硬件规格是google colab，cpu是Intel(R) Xeon(R) CPU @ 2.00GHz，GPU是Tesla T4。
从结果来看，在CPU上，TorchScript比pytorch eager快了3.5%，在GPU上，TorchScript比pytorch快了55.6%。
然后我们再用ResNet做一个测试。
import torchvision
import torch
from time import perf_counter
import numpy as npdef timer(f,*args): start = perf_counter()f(*args)return (1000 * (perf_counter() - start))# Pytorch cpu versionmodel_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 versionmodel_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 versionscript_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 versionscript_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 | 77.47 | 2.99 |
| TorchScript | 74.24 | 1.64 |
TorchScript相比PyTorch eager model，CPU性能提升4.2%，GPU性能提升45%。与Bert的结论一致。
总结
- 本文重点说明了Pytorch的eager模式和script模式，重点是script模式的TorchScript和Pytorch JIT
- 上一篇文章重点说明了eager模式的model转为script模式的TorchScript的两个api，
torch.jit.trace与torch.jit.script的区别，这是这一篇文章的基础，建议先阅读上一篇文章 - 使用Bert和ResNet两个网络进行了Pytorch eager model和TorchScript的CPU和GPU性能测试。结论在两个网络上一致，使用TorchScript在CPU上，相比PyTorch eager mode，会有4%左右的性能提升，在GPU上，会有50%左右的性能提升。
】)
