创建 TorchScript 模块¶
TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。PyTorch 有详细的文档说明如何实现这一点:https://pytorch.ac.cn/tutorials/beginner/Intro_to_TorchScript_tutorial.html。这里简要介绍一下关键背景信息和过程。
PyTorch 程序基于 Module
,Module
可用于组合更高级别的模块。Module
包含一个用于设置模块、参数和子模块的构造函数,以及一个描述模块调用时如何使用参数和子模块的前向函数 (forward function)。
例如,我们可以这样定义一个 LeNet 模块
1import torch.nn as nn
2import torch.nn.functional as F
3
4
5class LeNetFeatExtractor(nn.Module):
6 def __init__(self):
7 super(LeNetFeatExtractor, self).__init__()
8 self.conv1 = nn.Conv2d(1, 6, 3)
9 self.conv2 = nn.Conv2d(6, 16, 3)
10
11 def forward(self, x):
12 x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
13 x = F.max_pool2d(F.relu(self.conv2(x)), 2)
14 return x
15
16
17class LeNetClassifier(nn.Module):
18 def __init__(self):
19 super(LeNetClassifier, self).__init__()
20 self.fc1 = nn.Linear(16 * 6 * 6, 120)
21 self.fc2 = nn.Linear(120, 84)
22 self.fc3 = nn.Linear(84, 10)
23
24 def forward(self, x):
25 x = torch.flatten(x, 1)
26 x = F.relu(self.fc1(x))
27 x = F.relu(self.fc2(x))
28 x = self.fc3(x)
29 return x
30
31
32class LeNet(nn.Module):
33 def __init__(self):
34 super(LeNet, self).__init__()
35 self.feat = LeNetFeatExtractor()
36 self.classifier = LeNetClassifier()
37
38 def forward(self, x):
39 x = self.feat(x)
40 x = self.classifier(x)
41 return x
.
显然,您可能希望将如此简单的模型整合到一个模块中,但这里我们可以看到 PyTorch 的可组合性。
从 PyTorch Python 代码到 TorchScript 代码有两种途径:追踪 (Tracing) 和脚本化 (Scripting)。
追踪 (Tracing) 在模块被调用时跟踪执行路径并记录发生的情况。要追踪我们的 LeNet 模块实例,我们可以使用一个示例输入调用 torch.jit.trace
。
import torch
model = LeNet()
input_data = torch.empty([1, 1, 32, 32])
traced_model = torch.jit.trace(model, input_data)
脚本化 (Scripting) 实际上是使用编译器检查您的代码并生成等效的 TorchScript 程序。区别在于,由于追踪 (tracing) 跟踪的是模块的执行过程,它无法捕捉控制流等信息。而脚本化 (scripting) 通过分析 Python 代码,编译器可以包含这些组件。我们可以通过调用 torch.jit.script
在我们的 LeNet 模块上运行脚本编译器。
import torch
model = LeNet()
script_model = torch.jit.script(model)
选择哪种途径都有其原因,PyTorch 文档提供了如何选择的信息。从 Torch-TensorRT 的角度来看,对追踪 (traced) 模块的支持更好(即您的模块更有可能被编译),因为它不包含完整编程语言的所有复杂性,尽管两种途径都受支持。
在对模块进行脚本化 (scripting) 或追踪 (tracing) 后,您将获得一个 TorchScript 模块。它包含用于运行模块的代码和参数,存储在一个 Torch-TensorRT 可以使用的中间表示 (intermediate representation) 中。
以下是 LeNet 追踪 (traced) 模块的 IR 外观
graph(%self.1 : __torch__.___torch_mangle_10.LeNet,
%input.1 : Float(1, 1, 32, 32)):
%129 : __torch__.___torch_mangle_9.LeNetClassifier = prim::GetAttr[name="classifier"](%self.1)
%119 : __torch__.___torch_mangle_5.LeNetFeatExtractor = prim::GetAttr[name="feat"](%self.1)
%137 : Tensor = prim::CallMethod[name="forward"](%119, %input.1)
%138 : Tensor = prim::CallMethod[name="forward"](%129, %137)
return (%138)
以及 LeNet 脚本化 (scripted) 模块的 IR
graph(%self : __torch__.LeNet,
%x.1 : Tensor):
%2 : __torch__.LeNetFeatExtractor = prim::GetAttr[name="feat"](%self)
%x.3 : Tensor = prim::CallMethod[name="forward"](%2, %x.1) # x.py:38:12
%5 : __torch__.LeNetClassifier = prim::GetAttr[name="classifier"](%self)
%x.5 : Tensor = prim::CallMethod[name="forward"](%5, %x.3) # x.py:39:12
return (%x.5)
您可以看到 IR 保留了我们在 python 代码中的模块结构。
在 Python 中使用 TorchScript¶
TorchScript 模块的运行方式与普通 PyTorch 模块相同。您可以使用 forward
方法或直接调用模块(如 torch_script_module(in_tensor)
)来运行前向传播 (forward pass)。JIT 编译器将即时编译和优化模块,然后返回结果。
将 TorchScript 模块保存到磁盘¶
对于追踪 (traced) 或脚本化 (scripted) 模块,您可以使用以下命令将模块保存到磁盘
import torch
model = LeNet()
script_model = torch.jit.script(model)
script_model.save("lenet_scripted.ts")