创建 TorchScript 模块

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

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

例如，我们可以这样定义 LeNet 模块

import torch.nn as nn
import torch.nn.functional as F


class LeNetFeatExtractor(nn.Module):
    def __init__(self):
        super(LeNetFeatExtractor, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        return x


class LeNetClassifier(nn.Module):
    def __init__(self):
        super(LeNetClassifier, self).__init__()
        self.fc1 = nn.Linear(16 * 6 * 6, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.feat = LeNetFeatExtractor()
        self.classifier = LeNetClassifier()

    def forward(self, x):
        x = self.feat(x)
        x = self.classifier(x)
        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="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 模块 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) 来运行前向传递。JIT 编译器会即时编译和优化模块，然后返回结果。

将 TorchScript 模块保存到磁盘

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

import torch

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