PyTorch JIT和TorchScript

2 minute read

Published:

PyTorch支持两种模式:eager模式和script模式。eager模式主要用于模型的编写、训练和调试,script模式主要是针对部署的,其包含PytorchJIT和TorchScript(一种在 PyTorch 中执行高效的序列化代码格式)。

script模式使用torch.jit.tracetorch.jit.script创建一个PyTorch eager module的中间表示(intermediate representation, IR),IR 经过内部优化,并在运行时使用 PyTorch JIT 编译。PyTorch JIT 编译器使用运行时信息来优化 IR。该 IR 与 Python 运行时是解耦的。

PyTorch JIT(Just-In-Time Compilation)是 PyTorch 中的即时编译器。

  1. 它允许你将模型转化为 TorchScript 格式,从而提高模型的性能和部署效率。
  2. JIT 允许你在动态图和静态图之间无缝切换。你可以在 Python 中以动态图的方式构建和调试模型,然后将模型编译为 TorchScript 以进行优化和部署。
  3. JIT 允许你在不同的深度学习框架之间进行模型转换,例如将 PyTorch 模型转换为 ONNX 格式,从而可以在其他框架中运行。

TorchScript 是 PyTorch 提供的一种将模型序列化以便在其他环境中运行的机制。它将 PyTorch 模型编译成一种中间表示形式,可以在没有 Python 解释器的环境中运行。这使得模型可以在 C++ 等其他语言中运行,也可以在嵌入式设备等资源受限的环境中实现高效的推理。

以下是 TorchScript 的一些重要特性和用途:

  1. 静态图表示形式:TorchScript 是一种静态图表示形式,它在模型构建阶段对计算图进行编译和优化,而不是在运行时动态构建。这可以提高模型的执行效率。
  2. 模型导出:TorchScript 允许将 PyTorch 模型导出到一个独立的文件中,然后可以在没有 Python 环境的设备上运行。
  3. 跨平台部署:TorchScript 允许在不同的深度学习框架之间进行模型转换,例如将 PyTorch 模型转换为 ONNX 格式,从而可以在其他框架中运行。
  4. 模型优化和量化:通过 TorchScript,你可以使用各种技术(如量化)对模型进行优化,从而减小模型的内存占用和计算资源消耗。
  5. 融合和集成:TorchScript 可以帮助你将多个模型整合到一个整体流程中,从而提高系统的整体性能。
  6. 嵌入式设备:对于资源受限的嵌入式设备,TorchScript 可以帮助你优化模型以适应这些环境。

使用 TorchScript 可以将 PyTorch 模型变得更容易在生产环境中部署和集成。然而,它也可能需要你对模型进行一些修改以使其可以成功编译为 TorchScript。

总的来说,TorchScript 是一个强大的工具,特别是对于需要在不同环境中部署 PyTorch 模型的情况。通过将模型导出为 TorchScript,你可以实现更广泛的模型应用和部署。

一段话总结,为什么要用以及什么时候要用script模式呢?

  1. 可以脱离python GIL以及python runtime的限制来运行模型,比如通过LibTorch通过C++来运行模型。这样方便了模型部署,例如可以在IoT等平台上运行。例如这个tutorial,使用C++来运行pytorch的model。
  2. 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.tracetorch.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_counter

def 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 = 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]

# 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)
PyTorch171.2730.42
TorchScript165.2413.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 np

def timer(f,*args):   
    start = perf_counter()
    f(*args)
    return (1000 * (perf_counter() - start))
  
# Pytorch cpu version

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 version

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 version

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 version

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)
PyTorch77.472.99
TorchScript74.241.64

TorchScript相比PyTorch eager model,CPU性能提升4.2%,GPU性能提升45%。与Bert的结论一致。

总结

  1. 本文重点说明了Pytorch的eager模式和script模式,重点是script模式的TorchScript和Pytorch JIT
  2. 上一篇文章重点说明了eager模式的model转为script模式的TorchScript的两个api,torch.jit.tracetorch.jit.script的区别,这是这一篇文章的基础,建议先阅读上一篇文章
  3. 使用Bert和ResNet两个网络进行了Pytorch eager model和TorchScript的CPU和GPU性能测试。结论在两个网络上一致,使用TorchScript在CPU上,相比PyTorch eager mode,会有4%左右的性能提升,在GPU上,会有50%左右的性能提升。