Torch JIT加速推理
JIT
如果多次使用某一个正则表达式，建议先对其进行compile，再使用compile之后得到的对象来做正则匹配。而这个compile的过程就是JIT(即时编译)。
举个栗子，正则表达式匹配有两种方式
pat = re.compile(pattern)
result = pat.match(string)
或
result = re.match(pattern, string)
前者先编译，再匹配，如果多次使用该正则表达匹配，建议先编译，之后就只用调用transfer匹配了
torch.jit > TorchScript
PyTorch 的 1.0 版本发布的最核心的两个新特性就是 JIT 和 C++ API，这两个特性一起发布不是没有道理的，JIT 是 Python 和 C++ 的桥梁。开发者可以使用 Python 训练模型，然后通过 JIT 将模型转为语言无关的模块， 从而让 C++ 可以非常方便得调用，从此使用 Python 训练模型，使用 C++ 将模型部署到生产环境对 PyTorch 来说成为了一件很容易的事。而因为使用了 C++， 开发者几乎可以把 PyTorch 模型部署到任意平台和设备上：树莓派、iOS、Android 等。
TorchScript是PyTorch模型的一种中间表示(即常见的intermediate representation，IR)，是nn.Module的子类，基于这种模型格式可以在高性能的环境如C++中运行。TorchScript这种PyTorch模型表示格式可以由 TorchScript 编译器理解、编译和序列化而成。从根本上说，TorchScript 本身就是一种编程语言。它是使用 PyTorch API 的 Python 的子集。简单来说，TorchScript 软件栈可以将 Python 代码转换成 C++ 代码。TorchScript 软件栈包括两部分：TorchScript(Python)和 LibTorch(C++)。TorchScript负责将 Python 代码转成一个模型文件，LibTorch负责解析运行这个模型文件。后续会从代码层面介绍如何保存为TorchScript模型格式及其在Python和C++中如何加载使用
- 有torch.jit.trace、torch.jit.script
- torch.jit.trace适用于静态图模型，torch.jit.script适用于动态图模型（模型的结构和输入在运行时变换）
trace
trace 模式顾名思义就是跟踪模型的执行，然后记录执行过程中的路径。在使用 trace 模式时，需要构造一个符合要求的输入，然后使用 TorchScript tracer 运行一遍，记录整个运行过程。在 trace 模式中运行时，每执行一个算子，就会往当前的 graph 加入一个 node。所有代码执行完毕，每一步的操作就会以一个计算图里的某个节点的形式被保存下来。
值得一提的是，PyTorch 导出 ONNX 也是使用了这部分代码，所以理论上能够导出 ONNX 的模型也能够使用 trace 模式导出 TorchScript 格式模型。
在VCAP中，将onnx模型转为vaimlite（CPU/GPU)或bin+params(NPU），之后的输入形状必须固定，猜测之前是可以的
trace 模式有以下2点限制：
- 不能有 if-else 等控制流
- 只支持 Tensor 操作。不支持非Tensor 操作，如List、Tuple、Map 等容器操作
通过 trace 模式的特点可以看出，trace 模式通常用于深度模型的导出。常见的深度模型通常没有 if-else 控制流且通常没有非 Tensor 操作。
实例：
import torch
torch.manual_seed(42)
class TestTraceCell(torch.nn.Module):
def __init__(self):
super(TestTraceCell, self).__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return new_h, new_h
test_cell = TestTraceCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(test_cell, (x, h))
print(traced_cell)
print("trace module=",traced_cell(x, h))
script
script 模式和trace 模式的使用方式很接近， 但是实现原理却大相径庭。TorchScript 实现了一个完整的编译器以支持 script 模式。保存模型阶段对应着编译器的前端(语法分析、语义分析、类型检查、中间代码生成)。即，对code进行词法分析、语法分析、句法分析，形成AST树(抽象语法树)，最后再将AST树线性化。在保存模型时，TorchScript 编译器解析 Python 代码，并构建代码的AST。script 模式在的限制比较小，不仅支持 if-else 等控制流，还支持非 Tensor 操作，如 List、Tuple、Map 等容器操作。script相当于一个嵌入在Python/Pytorch的DSL，其语法只是pytorch语法的子集，这意味着存在一些op和语法script不支持，这样在编译的时候就会遇到问题。此外，script的编译优化方式更像是CPU上的传统编译优化，重点对于图进行硬件无关优化，并对IF、loop这样的statement进行优化。
import torch
torch.manual_seed(42)
class TestDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
class TestCell(torch.nn.Module):
def __init__(self, dg):
super(TestCell, self).__init__()
self.dg = dg
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
test_cell = TestCell(TestDecisionGate())
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(test_cell, (x, h))
print(traced_cell.code)
混合使用
import torch
torch.manual_seed(42)
# 无if-else，仅仅tensor op，用torch.jit.trace
class TestTraceCell(torch.nn.Module):
def __init__(self):
super(TestTraceCell, self).__init__()
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.linear(x) + h)
return new_h, new_h
# 有if-else，用script
class TestDecisionGate(torch.nn.Module):
def forward(self, x):
if x.sum() > 0:
return x
else:
return -x
# no if-else，only tensor，use trace
class TestCell(torch.nn.Module):
def __init__(self, dg):
super(TestCell, self).__init__()
self.dg = dg
self.linear = torch.nn.Linear(4, 4)
def forward(self, x, h):
new_h = torch.tanh(self.dg(self.linear(x)) + h)
return new_h, new_h
x, h = torch.rand(3, 4), torch.rand(3, 4)
scripted_gate = torch.jit.script(TestDecisionGate())
# for = if-else，use script，but member member can use trace
class TestRNNLoop(torch.nn.Module):
def __init__(self):
super(TestRNNLoop, self).__init__()
self.cell = torch.jit.trace(TestCell(scripted_gate), (x, h)) # TestCell本身用trace，但transfer的dg用script
def forward(self, xs):
h, y = torch.zeros(3, 4), torch.zeros(3, 4)
for i in range(xs.size(0)):
y, h = self.cell(xs[i], h)
return y, h
rnn_loop = torch.jit.script(TestRNNLoop())
print(rnn_loop.code)
在某些情况下，需要使用Tracing而不是Scripting，例如，module中有许多决策，这些决策是基于Python中固定值做出的，而我们又不希望在TorchScript中出现。在这种情况下，可以通过tracing来组成scripting：torch.jit.script将内联被traced module 的代码，而tracing将内联被scripted module 的代码。
torch.compile
torch.compile是PyTorch库中的一个函数，它主要用于将一段定义好的PyTorch脚本转换为可以高效运行的低级张量表达式（Tensor Expressions）。当你对一些计算过程进行了多次同样的操作，如前向传播、损失计算等，并且这个过程相对固定，torch.compile可以帮助优化这部分代码，使其在后续运行时速度更快，尤其是在分布式环境中或者模型部署阶段。
编译的过程会将动态图模式下的Python代码转化为静态图模式，这有助于消除不必要的函数调用开销和垃圾回收，提高性能。然而，需要注意的是，只有那些不会依赖于外部状态（如随机数生成或变量初始化）的纯函数才适合编译。
orient dynamic input
现在有一个输入的batch_size维度可能发生变化的pytorch模型，它适合用torch.compile进行编译码
对于一个输入batch_size可能会变化的PyTorch模型，通常不太适合直接使用torch.compile进行编译。因为torch.compile是针对静态图模式设计的，它期望所有的输入形状在编译时已经确定下来，以便生成高效的执行计划。如果模型的输入batch_size部分是动态的，那么每次传入的数据都会影响到计算路径，这种情况下静态图并不能很好地优化性能。
相反，如果你的模型有大量可重复使用的前向传播逻辑，但是batch_size的变化仅限于网络的部分层（例如卷积神经网络的池化层），你可以考虑把这部分逻辑封装成一个独立的模块，并在该模块内部处理batch_size的变化。然后，这个模块就可以作为静态图的一部分来编译，而其他部分继续保持动态图的灵活性。
所以，在这种情况下，更好的做法可能是拆分模型，只对那些可以提前确定计算流程的部分进行编译。
Questions
Q1:torch.jit.trace支持列表的索引操作吗
C知道说torch.jit.trace不支持，torch.jit.script支持，但本地试验的时候，都是支持的
Q2:torch.jit.trace如何接受多级input？
- 直接按照教程的拼接方法再用torch.jit.trace()会bloom
:List inputs to traced functions must have consistent element type. Found Tensor and Tuple[Tensor, Tensor]
discuss1
RuntimeError: Only tensors, lists, tuples of tensors, or dictionary of tensors can be output from traced functions
script
2024-10-14 22:03:51.730 | ERROR | actor.agent:reset:102 - could not get source code
2024-10-14 22:03:51.730 | ERROR | actor.agent:reset:103 - get_model error, try again...
Reference
- 模型推理加速系列 | 05: 推理加速格式TorchScript简介及其应用
- torch.jit.trace与torch.jit.script的区别
- torch.jit.trace Official Guide
鄙人孤陋寡闻，那有什么学问，还怕人偷不成?