• 文档 >
  • 创建 TorchScript 模块
快捷方式

创建 TorchScript 模块

TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。 PyTorch 有关于如何执行此操作的详细文档 https://pytorch.ac.cn/tutorials/beginner/Intro_to_TorchScript_tutorial.html 但简而言之,以下是关键背景信息和过程

PyTorch 程序基于 Module,这些模块可用于组成更高级别的模块。 Modules 包含一个用于设置模块、参数和子模块的构造函数,以及一个描述如何在调用模块时使用参数和子模块的正向函数。

例如,我们可以像这样定义一个 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.classifer = LeNetClassifier()
37
38    def forward(self, x):
39        x = self.feat(x)
40        x = self.classifer(x)
41        return x

.

显然,您可能希望将这样一个简单的模型合并到一个模块中,但我们可以在这里看到 PyTorch 的可组合性

从这里开始,有两条路径可以从 PyTorch Python 代码转换为 TorchScript 代码:跟踪和脚本。

跟踪遵循模块被调用时的执行路径并记录发生的事情。 要跟踪 LeNet 模块的实例,我们可以使用示例输入调用 torch.jit.trace

import torch

model = LeNet()
input_data = torch.empty([1, 1, 32, 32])
traced_model = torch.jit.trace(model, input_data)

脚本实际上是用编译器检查您的代码并生成等效的 TorchScript 程序。 区别在于,由于跟踪遵循模块的执行,因此它无法获取例如控制流。 通过使用 Python 代码,编译器可以包含这些组件。 我们可以通过调用 torch.jit.script 在 LeNet 模块上运行脚本编译器

import torch

model = LeNet()
script_model = torch.jit.script(model)

有理由使用一种或另一种路径,PyTorch 文档提供了有关如何选择的信息。 从 Torch-TensorRT 的角度来看,对跟踪模块有更好的支持(即您的模块更有可能编译),因为它不包括完整编程语言的所有复杂性,尽管两种路径都支持。

在编写或跟踪模块后,您将获得一个 TorchScript 模块。 它包含用于运行存储在 Torch-TensorRT 可以使用的中间表示中的模块的代码和参数。

以下是 LeNet 跟踪模块 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="classifer"](%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 脚本模块 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="classifer"](%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_scirpt_module(in_tensor) JIT 编译器将即时编译和优化模块,然后返回结果。

将 TorchScript 模块保存到磁盘

对于跟踪模块或脚本模块,您可以使用以下命令将模块保存到磁盘

import torch

model = LeNet()
script_model = torch.jit.script(model)
script_model.save("lenet_scripted.ts")

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源