pytorch JIT 和 TorchScript 详解

笔者自己对这个也只是一知半解，如果内容有出入，欢迎指出。后续给出的4个参考链接有很详细的说明，建议详看参考文献。

一、PyTorch 特性

我们都喜欢PyTorch的动态性和易用性。但在部署方面，这些品质不如性能和可移植性理想。为了提高pytorch的可移植性，引入了 TorchScript机制。

首先，探究一下 PyTorch 的生态系统，PyTorch 主要支持两种模式:

1. Eager mode: 它用于构建原型、训练和实验
2. Script mode：它主要关注在生产用例方面，包含 PyTorch JIT 和 TorchScript两个部分。

那么，想想我们为什么需要Script mode呢？

GIL是什么 详情参考博客园 python中的GIL详解
首先需要明确的一点是GIL并不是Python的特性，它是在实现Python解析器(CPython)时所引入的一个概念。就好比C++是一套语言（语法）标准，但是可以用不同的编译器来编译成可执行代码。有名的编译器例如GCC，INTEL C++，Visual C++等。Python也一样，同样一段代码可以通过CPython，PyPy，Psyco等不同的Python执行环境来执行。像其中的JPython就没有GIL。然而因为CPython是大部分环境下默认的Python执行环境。所以在很多人的概念里CPython就是Python，也就想当然的把GIL归结为Python语言的缺陷。所以

这里要先明确一点：GIL并不是Python的特性，Python完全可以不依赖于GIL。
那么CPython实现中的GIL又是什么呢？GIL全称Global Interpreter Lock为了避免误导，来看一下官方给出的解释：
In CPython, the global interpreter lock, or GIL, is a mutex that prevents multiple native threads from executing Python bytecodes at once. This lock is necessary mainly because CPython’s memory management is not thread-safe. (However, since the GIL exists, other features have grown to depend on the guarantees that it enforces.)

简单的说，就是为了摆脱 python 的 GIL 特性和 运行的python环境依赖。一个巧妙的解释如下：

1. 可移植性(Portability): 可移植到任何没有特定条件的运行环境中运行。
2. 性能(Performance)：PyTorch JIT 是一个为优化PyTorch的JIT编译器，它使用运行信息来优化TorchScript 模块。他可以自动优化 layer fusion，quantization, sparsification.

二、脚本模式(Script Mode)

script mode 通过 torch.jit.trace 或 torch.jit.script为PyTorch 的 eager modules 创建了一个 intermediate representation（IR，中间解释器），IR在内部进行了优化，并在运行时使用了PyTorch JIT编译。PyTorch JIT编译器使用运行时信息来优化IR。这个IR与Python运行时解耦。
脚本模式通过使用 PyTorch JIT 和 TorchScript 工作。

1. 什么是PyTorch JIT？

PyTorch JIT是为PyTorch程序优化的编译器。

• 它是一个轻量级的thread safe解释器
• 支持易于编写自定义转换
• 它不只是用于推理，因为它有自动差值支持

2. 什么是TorchScript?

TorchScript是Python语言的一个静态高性能子集，专门用于ML应用程序。它支持

• 复杂控制流
• 常见的数据结构
• 用户定义的类

3. 使用 Script mode

script mode 引入了 torch.jit.trace 或 torch.jit.script 实现。

（1） torch.jit.trace

torch.jit.trace 跟踪以数据实例和训练有素的模块作为输入。跟踪程序运行提供的模块并记录执行的张量操作。这个记录被转换成一个TorchScript模块。
示例详见 PyTorch JIT and TorchScript 的 Example 1: BERT 。

（2） torch.jit.script

torch.jit.script 允许你直接在TorchScript中编写代码。它更啰嗦，但更通用，稍加调整就可以支持大多数PyTorch模型。
示例详见 PyTorch JIT and TorchScript 的 Example 2: ResNet。

参考文献：
PyTorch JIT and TorchScript
官方文档 INTRODUCTION TO TORCHSCRIPT
google AI blog Eager Execution: An imperative, define-by-run interface to TensorFlow
博客园 python中的GIL详解