PyTorch JIT简介
PyTorch JIT用于编译用PyTorch编写的程序。它是针对PyTorch的一个完全优化的编译器解决方案。PyTorch JIT具有多个功能和特点，使其更加完美。其中包括具有在轻量级和安全的环境中进行解释的额外功能，可以很容易地编写自定义转换，并且支持推理，它还带有自动差异支持。
在本文中，我们将了解什么是PyTorch JIT，使用PyTorch JIT、JIT跟踪和JIT脚本，PyTorch JIT模块，一些与之相关的例子，最后，总结一下我们的发言。
什么是PyTorch JIT？
JIT编译器利用运行时信息对Torchscript的模块进行优化。使用JIT即Just In Time编译器在优化稀疏化、量化和层融合方面可以实现自动化。
脚本模式利用eager模块和torch.JIT.script及torch.JIT.trace模块来创建IR，即PyTorch的中间表示法。
使用PyTorch JIT
我们需要创建模型，将其转移到生产环境中，然后使用PyTorch JIT编译器对其进行编译和优化，使推理工作不费吹灰之力就能更快。
PyTorch使用两种模式处理生产和研究：急切和脚本模式。Eager用于培训、原型设计和实验，而脚本模式则用于生产。脚本模式包含两个独立的组件，即TorchScript和PyTorch JIT。脚本模式被证明对避免依赖python和python的GIL有很大帮助。此外，由于使用了PyTorch JIT这个优化的编译器，它还带来了可移植性和性能等特点。
JIT跟踪和JIT脚本
脚本模式的调用是通过使用torch.JIT.script或通过使用torch.JIT.trace来进行的。作为torch.JIT.trace的JIT trace接受包含我们训练好的、渴望的模型和数据实例的输入。JIT追踪器负责保持张量的操作记录，并帮助运行提供的模块。此外，我们将JIT追踪器的记录转换为火炬脚本的模块。
JIT脚本也被称为Torch.JIT.script，使我们直接在Torch脚本里面写代码。它是通用的，同时也是有点啰嗦的。经过一定的调整，它可以帮助我们支持大部分的PyTorch模型。我们只需要向JIT脚本提供模型实例，这与JIT跟踪模式完全不同。没有必要指定数据的样本。
PyTorch JIT模块
脚本模式利用eager模块和torch.JIT.script及torch.JIT.trace模块来创建IR，即PyTorch的中间表示法。PyTorch的JIT模块是JIT跟踪模块和JIT脚本模块。
JIT跟踪模块
如果存在急切的模型，Tracer可以重复使用编码器，并使用独占的Torch操作和张量来处理所有的程序。追踪器的缺点包括省略了python结构、数据结构和控制流。即使没有指定任何警告，它也会创建它们不忠实的表示法。只是为了验证PyTorch的模型结构是否被正确解析，它总是检查IR。
JIT脚本模块
它看起来与Python相似，并支持Python的结构，而且控制流被保留下来。此外，它还提供了对字典和列表的完整支持。然而，它不支持常量值，类型转换是必要的。当没有指定任何类型时，默认的类型是张量。
PyTorch JIT的例子
让我们考虑一个例子，这个例子是由google人工智能开发的，名为BERT，代表来自变形器的双向编码器表示法。我们将使用的库是由hugging face提供的。
我们将按照下面的步骤来实现这个例子--
- 将初始化标记器和BERT模型，并创建用于推理的样本数据。
- PyTorch模型将准备在GPU和CPU上进行推理。
- Torchscript模型应准备在GPU和CPU上进行推理。
from transformers
import numpy as numpyObj
import torch
import BertTokenizer, BertModel
from time import perf_counter
def sampleTimerFun(f,*args):
initiate = perf_counter()
f(*args)
return (1000 * (perf_counter() - initiate))
sampleTokenizerForScript = BertTokenizer.from_pretrained('bert-base-uncased', torchscript=True)
modelOfScript = BertModel.from_pretrained("bert-base-uncased", torchscript=True)
# Text in input is being tokenized
inputString = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_inputString = sampleTokenizerForScript.tokenize(inputString)
# One of the input token is being masked
indexOfMaskedToken = 8
tokenized_inputString[indexOfMaskedToken] = '[MASK]'
ListOfIndexedTokens = sampleTokenizerForScript.convert_tokens_to_ids(tokenized_inputString)
IdOfSegments = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1] # dummy object creation for input value
tensorOfTokens = torch.tensor([ListOfIndexedTokens])
tensorOfSegments = torch.tensor([IdOfSegments])
# 例子1.1 - 在CPU上演示BERT模型的例子
sampleNativeModel = BertModel.from_pretrained("bert-base-uncased")
numpyObj.mean([sampleTimerFun(sampleNativeModel,tensorOfTokens,tensorOfSegments) for _ in range(100)])
# 例子1.2 - 在GPU上演示BERT模型的例子
# 为了应用推理，用于采样的两个模型都应该存在于样本的GPU设备上。
sampleGPUOfNative = sampleNativeModel.cuda()
tensorOfTokens_gpu = tensorOfTokens.cuda()
tensorOfSegments_gpu = tensorOfSegments.cuda()
numpyObj.mean([sampleTimerFun(sampleGPUOfNative,tensorOfTokens_gpu,tensorOfSegments_gpu) for _ in range(100)])
# 例1.3 - CPU实现的torch.jit.trace
JITTraceModel = torch.jit.trace(modelOfScript, [tensorOfTokens, tensorOfSegments])
numpyObj.mean([sampleTimerFun(JITTraceModel,tensorOfTokens,tensorOfSegments) for _ in range(100)])
# 例1.4 - GPU实现的torch.jit.trace
JITTraceModel_gpu = torch.jit.trace(modelOfScript.cuda(), [tensorOfTokens.cuda(), tensorOfSegments.cuda()])
numpyObj.mean([sampleTimerFun(JITTraceModel_gpu,tensorOfTokens.cuda(),tensorOfSegments.cuda()) for _ in range(100)])
我们可以观察到，CPU上的运行时间几乎相同，但在GPU上，torchscript的运行时间被证明比PyTorch好。
在内部，torchcript创建了PyTorch模型的中间表示，可以在运行时通过使用及时编译器（JIT）进行优化编译。我们可以通过traced_model.code来看一下作为中间表示的IR。
总结
我们了解了脚本模式和PyTorch JIT。通过使用脚本模式，可以在生产环境中创建用于研究的模型，同时，它仍然在PyTorch的生态系统中。
