pytorch TorchScript 简介

本文介绍 TorchScript，它是 PyTorch 模型（nn.Module的子类）的中间表示，可以在 C++ 等高性能环境中运行。

在本文中，我们将介绍：

1. PyTorch 中模型构建的基础知识，包括：

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

2. 将 PyTorch 模块转换为 TorchScript（高性能部署运行时）的具体方法

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

我们希望在您完成本教程后，您将继续 学习后续教程 ，该教程将引导您了解从 C++ 实际调用 TorchScript 模型的示例。

import torch  # This is all you need to use both PyTorch and TorchScript!
print(torch.__version__)

1.12.1+cu102

1. PyTorch 模型构建基础

让我们定义一个简单的Module. Module是 PyTorch 中基本的组成单位。它包含：

1. 构造函数，在这里会做一些初始化
2. 一组Parameters和Modules子类。这些由构造函数初始化，并且可以在调用期间由模块使用。
3. 一个forward函数。这是调用模块时运行的代码。

让我们看一个小例子：

class MyCell(torch.nn.Module):
    # 构造函数
    def __init__(self):
        super(MyCell, self).__init__()
    
    def forward(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.9433, 0.6199, 0.7729, 0.4915],
        [0.9102, 0.7947, 0.9392, 0.9473],
        [0.8212, 0.7222, 0.6727, 0.5907]]), tensor([[0.9433, 0.6199, 0.7729, 0.4915],
        [0.9102, 0.7947, 0.9392, 0.9473],
        [0.8212, 0.7222, 0.6727, 0.5907]]))

所以我们现在有：

1. 创建了一个torch.nn.Module子类。
2. 定义了一个构造函数。构造函数做的不多，只是调用父类super的构造函数.
3. 定义了一个forward函数，它接受两个输入并返回两个输出。forward函数的实际内容并不重要，但它有点像一个假的RNN 单元，也就是说它是一个应用于循环的函数。

我们实例化了模块，并初始化了x和h，它们只是 3x4 的随机值矩阵。然后我们用my_cell(x, h)发起的调用， 这样就会在内部调用我们的forward函数。

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

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, 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

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

MyCell(
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[0.7952, 0.6080, 0.7205, 0.4583],
        [0.8384, 0.7931, 0.8368, 0.8723],
        [0.8859, 0.5437, 0.3791, 0.3023]], grad_fn=<TanhBackward0>), tensor([[0.7952, 0.6080, 0.7205, 0.4583],
        [0.8384, 0.7931, 0.8368, 0.8723],
        [0.8859, 0.5437, 0.3791, 0.3023]], grad_fn=<TanhBackward0>))

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

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

printon aModule将给出 Module的子类层次结构的可视化表示。在我们的示例中，我们可以看到我们的 Linear子类及其参数。

通过以这种方式组合Module，我们可以简洁易读地创建具有可重用组件的模型。

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

现在让我们演示一下所所谓的灵活性：

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.dg = MyDecisionGate()
        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

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.8930,  0.7665,  0.3622, -0.0867],
        [ 0.8816,  0.8680,  0.7857, -0.0420],
        [ 0.8523,  0.7855,  0.4425, -0.6193]], grad_fn=<TanhBackward0>), tensor([[ 0.8930,  0.7665,  0.3622, -0.0867],
        [ 0.8816,  0.8680,  0.7857, -0.0420],
        [ 0.8523,  0.7855,  0.4425, -0.6193]], grad_fn=<TanhBackward0>))

我们再次重新定义了 MyCell 类，但在这里我们定义了 MyDecisionGate. 该模块利用控制流。控制流由循环和if语句之类的东西组成。

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

2. TorchScript 基础知识

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

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

2.1 Tracing Modules

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, 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

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.6235, 0.5857, 0.1536, 0.8709],
        [0.7608, 0.6827, 0.0668, 0.8764],
        [0.7824, 0.7744, 0.4454, 0.8038]], grad_fn=<TanhBackward0>), tensor([[0.6235, 0.5857, 0.1536, 0.8709],
        [0.7608, 0.6827, 0.0668, 0.8764],
        [0.7824, 0.7744, 0.4454, 0.8038]], grad_fn=<TanhBackward0>))

我们把代码回退了一个版本，使用了MyCell类的第二个版本。和以前一样，我们已经实例化了它，但是这一次，我们调用 torch.jit.trace了 ，传入了Module，并传入了网络可能会看到的示例输入。

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

TorchScript 将模型定义记录在中间表示（或 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/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %12 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::add(%20, %h, %11) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %13 : Float(3, 4, strides=[4, 1], requires_grad=1, device=cpu) = aten::tanh(%12) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188: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)

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

1. TorchScript 代码可以在它自己的解释器中调用，它基本上是一个受限的 Python 解释器。该解释器不获取全局解释器锁，因此可以在同一个实例上同时处理许多请求。
2. 这种格式允许我们将整个模型保存到磁盘并将其加载到另一个环境中，例如在用 Python 以外的语言编写的服务器中
3. TorchScript 为我们提供了一种表示，我们可以在其中对代码进行编译器优化以提供更有效的执行
4. TorchScript 允许我们与许多后端/设备运行时进行交互，这些运行时需要比单个操作更广泛的程序视图。 TorchScript allows us to interface with many backend/device runtimes that require a broader view of the program than individual operators.

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

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

(tensor([[0.6235, 0.5857, 0.1536, 0.8709],
        [0.7608, 0.6827, 0.0668, 0.8764],
        [0.7824, 0.7744, 0.4454, 0.8038]], grad_fn=<TanhBackward0>), tensor([[0.6235, 0.5857, 0.1536, 0.8709],
        [0.7608, 0.6827, 0.0668, 0.8764],
        [0.7824, 0.7744, 0.4454, 0.8038]], grad_fn=<TanhBackward0>))
(tensor([[0.6235, 0.5857, 0.1536, 0.8709],
        [0.7608, 0.6827, 0.0668, 0.8764],
        [0.7824, 0.7744, 0.4454, 0.8038]],
       grad_fn=<DifferentiableGraphBackward>), tensor([[0.6235, 0.5857, 0.1536, 0.8709],
        [0.7608, 0.6827, 0.0668, 0.8764],
        [0.7824, 0.7744, 0.4454, 0.8038]],
       grad_fn=<DifferentiableGraphBackward>))

3. 使用脚本转换模块

我们使用模块的第二个版本是有原因的，而不是带有控制流操作的子模块的那个。现在让我们仔细审查一下：

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self, dg):
        super(MyCell, 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

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

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

/var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:260: 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) -> Tensor:
  return torch.neg(argument_1)

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)
  _1 = torch.tanh(_0)
  return (_1, _1)

查看.code输出，我们可以看到if-else分支在代码中已经找不到了！为什么？跟踪完全按照我们所说的去做：运行代码，记录发生的操作并构造一个 ScriptModule 来完成这些工作。不幸的是，控制流之类的东西（比如 if 判断）被删除了。

我们如何在 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)):   # <<<<<<<<<===== 在这里看到了if语句
    _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)
traced_cell(x, h)

(tensor([[0.6409, 0.8655, 0.4093, 0.1051],
        [0.2452, 0.6884, 0.5398, 0.6191],
        [0.4787, 0.5705, 0.8379, 0.5657]], grad_fn=<TanhBackward0>), tensor([[0.6409, 0.8655, 0.4093, 0.1051],
        [0.2452, 0.6884, 0.5398, 0.6191],
        [0.4787, 0.5705, 0.8379, 0.5657]], grad_fn=<TanhBackward0>))

3.1 混合脚本和跟踪

某些情况需要使用跟踪而不是脚本（例如，一个模块有许多基于常量 Python 值做出的架构决策，我们不想出现在 TorchScript 中）。在这种情况下，脚本可以与跟踪组合：torch.jit.script将内联跟踪模块的代码，而跟踪将内联脚本模块的代码。

第一种情况的示例：

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

    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(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)

第二种情况的一个例子：

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

    def forward(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)

这样，当情况需要它们中的每一个并一起使用时，可以使用脚本和跟踪。

4.保存和加载模型

我们提供 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 执行。

原文地址