TorchScript 简介

作者: James Reed (jamesreed@fb.com), Michael Suo (suo@fb.com), rev2

TorchScript 已不再处于积极开发阶段。

本教程是对 TorchScript 的入门介绍，TorchScript 是 PyTorch 模型（nn.Module 的子类）的中间表示，可以在高性能环境（如 C++）中运行。

在本教程中，我们将涵盖：

PyTorch 中模型编写的基础知识，包括：

模块
定义 forward 函数
将模块组合成模块层次结构

将 PyTorch 模块转换为 TorchScript 的具体方法，TorchScript 是我们高性能的部署运行时

跟踪现有模块
使用脚本直接编译模块
如何结合两种方法
保存和加载 TorchScript 模块

我们希望您在完成本教程后，继续学习后续教程，该教程将引导您通过一个实际的示例，展示如何从 C++ 调用 TorchScript 模型。

importtorch  # This is all you need to use both PyTorch and TorchScript!
print(torch.__version__)
torch.manual_seed(191009)  # set the seed for reproducibility

2.6.0+cu124

<torch._C.Generator object at 0x7f3682196390>

PyTorch 模型编写基础

让我们从定义一个简单的 Module 开始。Module 是 PyTorch 中的基本组成单元。它包含：

构造函数，用于为模块的调用做准备
一组 Parameters 和子模块。这些由构造函数初始化，并可在模块调用期间使用。
forward 函数。这是模块被调用时执行的代码。

让我们来看一个小例子：

classMyCell(torch.nn.Module):
    def__init__(self):
        super(MyCell, self).__init__()

    defforward(self, x, h):
        new_h = torch.tanh(x + h)
        return new_h, new_h

my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))

(tensor([[0.8219, 0.8990, 0.6670, 0.8277],
        [0.5176, 0.4017, 0.8545, 0.7336],
        [0.6013, 0.6992, 0.2618, 0.6668]]), tensor([[0.8219, 0.8990, 0.6670, 0.8277],
        [0.5176, 0.4017, 0.8545, 0.7336],
        [0.6013, 0.6992, 0.2618, 0.6668]]))

因此，我们已经：

创建了一个继承自 torch.nn.Module 的类。
定义了一个构造函数。构造函数并没有做太多事情，只是调用了 super 的构造函数。
定义了一个 forward 函数，它接收两个输入并返回两个输出。forward 函数的具体内容并不重要，但它类似于一个假的 RNN 单元——也就是说，它是一个在循环中应用的函数。

我们实例化了该模块，并创建了 x 和 h，它们只是 3x4 的随机值矩阵。然后我们通过 my_cell(x, h) 调用了该单元。这反过来又调用了我们的 forward 函数。

让我们来做一些更有趣的事情：

classMyCell(torch.nn.Module):
    def__init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    defforward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

MyCell(
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.8573,  0.6190,  0.5774,  0.7869],
        [ 0.3326,  0.0530,  0.0702,  0.8114],
        [ 0.7818, -0.0506,  0.4039,  0.7967]], grad_fn=<TanhBackward0>), tensor([[ 0.8573,  0.6190,  0.5774,  0.7869],
        [ 0.3326,  0.0530,  0.0702,  0.8114],
        [ 0.7818, -0.0506,  0.4039,  0.7967]], grad_fn=<TanhBackward0>))

我们已经重新定义了我们的模块 MyCell，但这次我们添加了一个 self.linear 属性，并在 forward 函数中调用了 self.linear。

这里到底发生了什么？torch.nn.Linear 是 PyTorch 标准库中的一个 Module。就像 MyCell 一样，它可以通过调用语法来使用。我们正在构建一个 Module 的层次结构。

对 Module 使用 print 会显示出 Module 子类层次结构的可视化表示。在我们的例子中，我们可以看到 Linear 子类及其参数。

通过这种方式组合 Module，我们可以使用可重用的组件简洁且可读地编写模型。

你可能已经注意到输出中的 grad_fn。这是 PyTorch 自动微分方法（称为 autograd）的一个细节。简而言之，这个系统允许我们通过潜在复杂的程序计算导数。这种设计在模型编写中提供了极大的灵活性。

现在让我们来探讨这种灵活性：

classMyDecisionGate(torch.nn.Module):
    defforward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

classMyCell(torch.nn.Module):
    def__init__(self):
        super(MyCell, self).__init__()
        self.dg = MyDecisionGate()
        self.linear = torch.nn.Linear(4, 4)

    defforward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

MyCell(
  (dg): MyDecisionGate()
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.8346,  0.5931,  0.2097,  0.8232],
        [ 0.2340, -0.1254,  0.2679,  0.8064],
        [ 0.6231,  0.1494, -0.3110,  0.7865]], grad_fn=<TanhBackward0>), tensor([[ 0.8346,  0.5931,  0.2097,  0.8232],
        [ 0.2340, -0.1254,  0.2679,  0.8064],
        [ 0.6231,  0.1494, -0.3110,  0.7865]], grad_fn=<TanhBackward0>))

我们再次重新定义了 MyCell 类，但在这里我们定义了 MyDecisionGate。这个模块利用了控制流。控制流包括循环和 if 语句等结构。

许多框架采用的方法是，在给定完整程序表示的情况下计算符号导数。然而，在 PyTorch 中，我们使用梯度磁带。我们记录操作的发生顺序，并在计算导数时反向回放。这样，框架就不必为语言中的所有结构显式定义导数。

自动求导的工作原理

TorchScript 基础

现在让我们以正在运行的示例为例，看看如何应用 TorchScript。

简而言之，TorchScript 提供了工具来捕获模型的定义，即使考虑到 PyTorch 的灵活和动态特性。让我们从研究所谓的 tracing 开始。

追踪 `Modules`

classMyCell(torch.nn.Module):
    def__init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    defforward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

MyCell(
  original_name=MyCell
  (linear): Linear(original_name=Linear)
)

(tensor([[-0.2541,  0.2460,  0.2297,  0.1014],
        [-0.2329, -0.2911,  0.5641,  0.5015],
        [ 0.1688,  0.2252,  0.7251,  0.2530]], grad_fn=<TanhBackward0>), tensor([[-0.2541,  0.2460,  0.2297,  0.1014],
        [-0.2329, -0.2911,  0.5641,  0.5015],
        [ 0.1688,  0.2252,  0.7251,  0.2530]], grad_fn=<TanhBackward0>))

我们稍微回溯了一下，采用了第二版的 MyCell 类。和之前一样，我们已经实例化了它，但这次我们调用了 torch.jit.trace，传入了 Module，并传入了网络可能会看到的示例输入。

这到底做了什么？它调用了 Module，记录了 Module 运行时发生的操作，并创建了一个 torch.jit.ScriptModule 的实例（其中 TracedModule 就是一个实例）。

TorchScript 将其定义记录在中间表示（Intermediate Representation，简称 IR）中，在深度学习中通常称为图。我们可以通过 .graph 属性来检查这个图：

print(traced_cell.graph)

graph(%self.1 : __torch__.MyCell,
      %x : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu),
      %h : Float(3, 4, strides=[4, 1], requires_grad=0, device=cpu)):
  %linear : __torch__.torch.nn.modules.linear.Linear = prim::GetAttr[name="linear"](%self.1)
  %20 : Tensor = prim::CallMethod[name="forward"](%linear, %x)
  %11 : int = prim::Constant[value=1]() # /var/lib/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
  %12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /var/lib/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
  %13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /var/lib/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:191:0
  %14 : (Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu), Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu)) = prim::TupleConstruct(%13, %13)
  return (%14)

然而，这是一种非常底层的表示形式，图中包含的大部分信息对最终用户来说并无用处。相反，我们可以使用 .code 属性来提供代码的 Python 语法解释：

print(traced_cell.code)

def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  linear = self.linear
  _0 = torch.tanh(torch.add((linear).forward(x, ), h))
  return (_0, _0)

那么，为什么我们要做这一切呢？有以下几个原因：

TorchScript 代码可以在其自身的解释器中调用，该解释器基本上是一个受限的 Python 解释器。这个解释器不会获取全局解释器锁 (GIL)，因此可以在同一个实例上同时处理多个请求。
这种格式允许我们将整个模型保存到磁盘，并加载到另一个环境中，例如用非 Python 语言编写的服务器中。
TorchScript 提供了一种表示形式，使我们能够对代码进行编译器优化，从而实现更高效的执行。
TorchScript 允许我们与许多后端/设备运行时进行交互，这些运行时需要对程序有更广泛的视图，而不仅仅是单个操作符。

我们可以看到，调用 traced_cell 产生的结果与 Python 模块相同：

print(my_cell(x, h))
print(traced_cell(x, h))

(tensor([[-0.2541,  0.2460,  0.2297,  0.1014],
        [-0.2329, -0.2911,  0.5641,  0.5015],
        [ 0.1688,  0.2252,  0.7251,  0.2530]], grad_fn=<TanhBackward0>), tensor([[-0.2541,  0.2460,  0.2297,  0.1014],
        [-0.2329, -0.2911,  0.5641,  0.5015],
        [ 0.1688,  0.2252,  0.7251,  0.2530]], grad_fn=<TanhBackward0>))
(tensor([[-0.2541,  0.2460,  0.2297,  0.1014],
        [-0.2329, -0.2911,  0.5641,  0.5015],
        [ 0.1688,  0.2252,  0.7251,  0.2530]], grad_fn=<TanhBackward0>), tensor([[-0.2541,  0.2460,  0.2297,  0.1014],
        [-0.2329, -0.2911,  0.5641,  0.5015],
        [ 0.1688,  0.2252,  0.7251,  0.2530]], grad_fn=<TanhBackward0>))

使用脚本进行模块转换

我们之所以使用了模块的第二版本，而不是包含控制流密集子模块的版本，是有原因的。现在让我们来探讨一下：

classMyDecisionGate(torch.nn.Module):
    defforward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

classMyCell(torch.nn.Module):
    def__init__(self, dg):
        super(MyCell, self).__init__()
        self.dg = dg
        self.linear = torch.nn.Linear(4, 4)

    defforward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))

print(traced_cell.dg.code)
print(traced_cell.code)

/var/lib/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:263: TracerWarning:

Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!

def forward(self,
    argument_1: Tensor) -> NoneType:
  return None

def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  dg = self.dg
  linear = self.linear
  _0 = (linear).forward(x, )
  _1 = (dg).forward(_0, )
  _2 = torch.tanh(torch.add(_0, h))
  return (_2, _2)

查看 .code 输出，我们可以看到 if-else 分支已经无处可寻！为什么？追踪功能正如我们所说的那样：运行代码，记录实际发生的操作，并构建一个完全按照这些操作执行的 ScriptModule。不幸的是，像控制流这样的结构会被抹去。

我们如何在 TorchScript 中准确地表示这个模块呢？我们提供了一个脚本编译器，它直接分析您的 Python 源代码并将其转换为 TorchScript。让我们使用脚本编译器来转换 MyDecisionGate：

scripted_gate = torch.jit.script(MyDecisionGate())

my_cell = MyCell(scripted_gate)
scripted_cell = torch.jit.script(my_cell)

print(scripted_gate.code)
print(scripted_cell.code)

def forward(self,
    x: Tensor) -> Tensor:
  if bool(torch.gt(torch.sum(x), 0)):
    _0 = x
  else:
    _0 = torch.neg(x)
  return _0

def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  dg = self.dg
  linear = self.linear
  _0 = torch.add((dg).forward((linear).forward(x, ), ), h)
  new_h = torch.tanh(_0)
  return (new_h, new_h)

太棒了！我们现在已经准确地将程序的行为转换为了 TorchScript。接下来让我们尝试运行这个程序：

# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
print(scripted_cell(x, h))

(tensor([[ 0.5679,  0.5762,  0.2506, -0.0734],
        [ 0.5228,  0.7122,  0.6985, -0.0656],
        [ 0.6187,  0.4487,  0.7456, -0.0238]], grad_fn=<TanhBackward0>), tensor([[ 0.5679,  0.5762,  0.2506, -0.0734],
        [ 0.5228,  0.7122,  0.6985, -0.0656],
        [ 0.6187,  0.4487,  0.7456, -0.0238]], grad_fn=<TanhBackward0>))

脚本与追踪结合使用

在某些情况下，使用 tracing 比 scripting 更为合适（例如，当某个模块的许多架构决策基于不希望出现在 TorchScript 中的 Python 常量值时）。在这种情况下，scripting 可以与 tracing 结合使用：torch.jit.script 会将 traced 模块的代码内联，而 tracing 则会将 scripted 模块的代码内联。

以下是第一种情况的示例：

classMyRNNLoop(torch.nn.Module):
    def__init__(self):
        super(MyRNNLoop, self).__init__()
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))

    defforward(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(MyRNNLoop())
print(rnn_loop.code)

def forward(self,
    xs: Tensor) -> Tuple[Tensor, Tensor]:
  h = torch.zeros([3, 4])
  y = torch.zeros([3, 4])
  y0 = y
  h0 = h
  for i in range(torch.size(xs, 0)):
    cell = self.cell
    _0 = (cell).forward(torch.select(xs, 0, i), h0, )
    y1, h1, = _0
    y0, h0 = y1, h1
  return (y0, h0)

第二种情况的示例:

classWrapRNN(torch.nn.Module):
    def__init__(self):
        super(WrapRNN, self).__init__()
        self.loop = torch.jit.script(MyRNNLoop())

    defforward(self, xs):
        y, h = self.loop(xs)
        return torch.relu(y)

traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)

def forward(self,
    xs: Tensor) -> Tensor:
  loop = self.loop
  _0, y, = (loop).forward(xs, )
  return torch.relu(y)

这样，在需要时，脚本和追踪可以各自使用，也可以结合使用。

模型的保存与加载

我们提供了API，用于以归档格式将TorchScript模块保存到磁盘或从磁盘加载。这种格式包括代码、参数、属性和调试信息，这意味着该归档文件是模型的一个独立表示，可以在完全独立的进程中加载。让我们来保存并加载我们封装的RNN模块：

traced.save('wrapped_rnn.pt')

loaded = torch.jit.load('wrapped_rnn.pt')

print(loaded)
print(loaded.code)

RecursiveScriptModule(
  original_name=WrapRNN
  (loop): RecursiveScriptModule(
    original_name=MyRNNLoop
    (cell): RecursiveScriptModule(
      original_name=MyCell
      (dg): RecursiveScriptModule(original_name=MyDecisionGate)
      (linear): RecursiveScriptModule(original_name=Linear)
    )
  )
)
def forward(self,
    xs: Tensor) -> Tensor:
  loop = self.loop
  _0, y, = (loop).forward(xs, )
  return torch.relu(y)

如您所见，序列化保留了模块层次结构以及我们一直在检查的代码。该模型还可以加载，例如，到 C++ 中，以实现无 Python 的执行。

延伸阅读

我们已经完成了教程！如需更深入的演示，请查看 NeurIPS 演示，了解如何使用 TorchScript 转换机器翻译模型：https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ

下载 Python 源代码: Intro_to_TorchScript_tutorial.py

下载 Jupyter 笔记本: Intro_to_TorchScript_tutorial.ipynb