介绍TorchScript

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

译者：松鼠

校验：松鼠

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

在本教程中，我们将介绍：

  1. PyTorch中的模型基础创建，包括：

    • 模块
    • 定义forward功能
    • 构成模块组成模块的层次结构

  2. 将PyTorch模块转换为TorchScript(我们的高性能部署运行时）的特定方法

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

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

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

输出：

  1.  1.2.0

PyTorch模型制作的基础

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

  1. 构造函数，为调用准备模块
  2. 一组Parameters和Modules。这些由构造函数初始化，并且可以在调用期间由模块使用。
  3. forward功能。这是调用模块时运行的代码。

让我们来看看一个小例子：

  1.     class MyCell(torch.nn.Module):
  2.         def __init__(self):
  3.             super(MyCell, self).__init__()
  5.         def forward(self, x, h):
  6.             new_h = torch.tanh(x + h)
  7.             return new_h, new_h
  9.     my_cell = MyCell()
  10.     x = torch.rand(3, 4)
  11.     h = torch.rand(3, 4)
  12.     print(my_cell(x, h))

Out:

  1.     (tensor([[0.7853, 0.8882, 0.7137, 0.3746],
  2.             [0.5265, 0.8508, 0.1487, 0.9144],
  3.             [0.7057, 0.8217, 0.9575, 0.6132]]), tensor([[0.7853, 0.8882, 0.7137, 0.3746],
  4.             [0.5265, 0.8508, 0.1487, 0.9144],
  5.             [0.7057, 0.8217, 0.9575, 0.6132]]))

因此，我们已经：

  1. 创建子类的类torch.nn.Module。
  2. 定义构造函数。构造函数没有做太多事情，只是调用的构造函数super。
  3. 定义了一个forward函数，该函数接受两个输入并返回两个输出。该forward函数的实际内容并不是很重要，但是它是一种伪造的RNN单元即，该函数应用于循环。

我们实例化了模块，并制作了x和y，它们只是3x4随机值矩阵。然后，我们使用调用单元,调用我们的forward函数my_cell(x, h)

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

  2.     class MyCell(torch.nn.Module):
  3.         def __init__(self):
  4.             super(MyCell, self).__init__()
  5.             self.linear = torch.nn.Linear(4, 4)
  7.         def forward(self, x, h):
  8.             new_h = torch.tanh(self.linear(x) + h)
  9.             return new_h, new_h
  11.     my_cell = MyCell()
  12.     print(my_cell)
  13.     print(my_cell(x, h))

Out:

  2.     MyCell(
  3.       (linear): Linear(in_features=4, out_features=4, bias=True)
  4.     )
  5.     (tensor([[0.7619, 0.7761, 0.7476, 0.0897],
  6.             [0.6886, 0.4990, 0.4495, 0.2021],
  7.             [0.5849, 0.5874, 0.9256, 0.0460]], grad_fn=<TanhBackward>), tensor([[0.7619, 0.7761, 0.7476, 0.0897],
  8.             [0.6886, 0.4990, 0.4495, 0.2021],
  9.             [0.5849, 0.5874, 0.9256, 0.0460]], grad_fn=<TanhBackward>))

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

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

print上的Module会直观地表示 Module的子类层次结构。在我们的示例中，我们可以看到 Linear子类及其参数。

通过这种方式构成Module们，我们可以简洁而易读地编写具有可复用组件的模型。

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

现在让我们检查一下灵活性：

  2.     class MyDecisionGate(torch.nn.Module):
  3.       def forward(self, x):
  4.         if x.sum() > 0:
  5.           return x
  6.         else:
  7.           return -x
  9.     class MyCell(torch.nn.Module):
  10.         def __init__(self):
  11.             super(MyCell, self).__init__()
  12.             self.dg = MyDecisionGate()
  13.             self.linear = torch.nn.Linear(4, 4)
  15.         def forward(self, x, h):
  16.             new_h = torch.tanh(self.dg(self.linear(x)) + h)
  17.             return new_h, new_h
  19.     my_cell = MyCell()
  20.     print(my_cell)
  21.     print(my_cell(x, h))

Out:

  2.     MyCell(
  3.       (dg): MyDecisionGate()
  4.       (linear): Linear(in_features=4, out_features=4, bias=True)
  5.     )
  6.     (tensor([[ 0.9077,  0.5939,  0.6809,  0.0994],
  7.             [ 0.7583,  0.7180,  0.0790,  0.6733],
  8.             [ 0.9102, -0.0368,  0.8246, -0.3256]], grad_fn=<TanhBackward>), tensor([[ 0.9077,  0.5939,  0.6809,  0.0994],
  9.             [ 0.7583,  0.7180,  0.0790,  0.6733],
  10.             [ 0.9102, -0.0368,  0.8246, -0.3256]], grad_fn=<TanhBackward>))

我们再次重新定义了MyCell类，但是在这里我们定义了 MyDecisionGate。该模块利用控制流程。控制流包括循环和if-statements之类的东西。

给定完整的程序表示形式，许多框架都采用计算符号派生的方法。但是，在PyTorch中，我们使用渐变色带。我们记录操作发生时的操作，并在计算衍生产品时向后回放。这样，框架不必为语言中的所有构造显式定义派生类。

How autograd works

autograd的工作原理

TorchScript的基础

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

简而言之，即使PyTorch具有灵活和动态的特性，TorchScript也提供了捕获模型定义的工具。让我们开始研究所谓的跟踪。

追踪Modules

  2.     class MyCell(torch.nn.Module):
  3.         def __init__(self):
  4.             super(MyCell, self).__init__()
  5.             self.linear = torch.nn.Linear(4, 4)
  7.         def forward(self, x, h):
  8.             new_h = torch.tanh(self.linear(x) + h)
  9.             return new_h, new_h
  11.     my_cell = MyCell()
  12.     x, h = torch.rand(3, 4), torch.rand(3, 4)
  13.     traced_cell = torch.jit.trace(my_cell, (x, h))
  14.     print(traced_cell)
  15.     traced_cell(x, h)

Out:

  1.     TracedModule[MyCell](
  2.       (linear): TracedModule[Linear]()
  3.     )

我们来看看之前的例子。和以前一样，我们实例化了它，但是这次，我们使用torch.jit.trace方法调用了Module，并传入，然后传入了网络可能的示例输入。

这到底是做什么的？它已调用Module，记录了Module运行时发生的操作，并创建了torch.jit.ScriptModule(TracedModule的实例）

TorchScript将其定义记录在中间表示(或IR）中，在深度学习中通常称为图形。我们可以检查具有以下.graph属性的图形：

  1.     print(traced_cell.graph)

Out:

  3.     graph(%self : ClassType<MyCell>,
  4.           %input : Float(3, 4),
  5.           %h : Float(3, 4)):
  6.       %1 : ClassType<Linear> = prim::GetAttr[name="linear"](%self)
  7.       %weight : Tensor = prim::GetAttr[name="weight"](%1)
  8.       %bias : Tensor = prim::GetAttr[name="bias"](%1)
  9.       %6 : Float(4!, 4!) = aten::t(%weight), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  10.       %7 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  11.       %8 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  12.       %9 : Float(3, 4) = aten::addmm(%bias, %input, %6, %7, %8), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1369:0
  13.       %10 : int = prim::Constant[value=1](), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  14.       %11 : Float(3, 4) = aten::add(%9, %h, %10), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  15.       %12 : Float(3, 4) = aten::tanh(%11), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  16.       %13 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%12, %12)
  17.       return (%13)

但是，这是一个非常低级的表示形式，图中包含的大多数信息对最终用户没有用。相反，我们可以使用.code属性为代码提供Python语法的解释：

  1.     print(traced_cell.code)

Out:

  1.     def forward(self,
  2.         input: Tensor,
  3.         h: Tensor) -> Tuple[Tensor, Tensor]:
  4.       _0 = self.linear
  5.       weight = _0.weight
  6.       bias = _0.bias
  7.       _1 = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
  8.       _2 = torch.tanh(torch.add(_1, h, alpha=1))
  9.       return (_2, _2)

那么为什么我们要做所有这些呢？有以下几个原因：

  1. TorchScript代码可以在其自己的解释器中调用，该解释器基本上是受限制的Python解释器。该解释器不获取全局解释器锁定，因此可以在同一实例上同时处理许多请求。
  2. 这种格式使我们可以将整个模型保存到磁盘上，并将其加载到另一个环境中，例如在以Python以外的语言编写的服务器中
  3. TorchScript为我们提供了一种表示形式，其中我们可以对代码进行编译器优化以提供更有效的执行
  4. TorchScript允许我们与许多后端/设备运行时进行接口，这些运行时比单个操作员需要更广泛的程序视图。

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

  2.     print(my_cell(x, h))
  3.     print(traced_cell(x, h))

Out:

  1.     (tensor([[ 0.0294,  0.2921,  0.5171,  0.2689],
  2.             [ 0.5859,  0.8311,  0.2553,  0.8026],
  3.             [-0.4138,  0.7641,  0.4251,  0.7217]], grad_fn=<TanhBackward>), tensor([[ 0.0294,  0.2921,  0.5171,  0.2689],
  4.             [ 0.5859,  0.8311,  0.2553,  0.8026],
  5.             [-0.4138,  0.7641,  0.4251,  0.7217]], grad_fn=<TanhBackward>))
  6.     (tensor([[ 0.0294,  0.2921,  0.5171,  0.2689],
  7.             [ 0.5859,  0.8311,  0.2553,  0.8026],
  8.             [-0.4138,  0.7641,  0.4251,  0.7217]],
  9.            grad_fn=<DifferentiableGraphBackward>), tensor([[ 0.0294,  0.2921,  0.5171,  0.2689],
  10.             [ 0.5859,  0.8311,  0.2553,  0.8026],
  11.             [-0.4138,  0.7641,  0.4251,  0.7217]],
  12.            grad_fn=<DifferentiableGraphBackward>))

使用脚本来转换模块

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

  1.     class MyDecisionGate(torch.nn.Module):
  2.       def forward(self, x):
  3.         if x.sum() > 0:
  4.           return x
  5.         else:
  6.           return -x
  8.     class MyCell(torch.nn.Module):
  9.         def __init__(self, dg):
  10.             super(MyCell, self).__init__()
  11.             self.dg = dg
  12.             self.linear = torch.nn.Linear(4, 4)
  14.         def forward(self, x, h):
  15.             new_h = torch.tanh(self.dg(self.linear(x)) + h)
  16.             return new_h, new_h
  18.     my_cell = MyCell(MyDecisionGate())
  19.     traced_cell = torch.jit.trace(my_cell, (x, h))
  20.     print(traced_cell.code)

Out:

  3.     def forward(self,
  4.         input: Tensor,
  5.         h: Tensor) -> Tuple[Tensor, Tensor]:
  6.       _0 = self.linear
  7.       weight = _0.weight
  8.       bias = _0.bias
  9.       x = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
  10.       _1 = torch.tanh(torch.add(torch.neg(x), h, alpha=1))
  11.       return (_1, _1)

查看.code输出，可以看到if-else找不到分支！为什么？跟踪完全按照我们所说的去做：运行代码，记录发生的操作，并构造一个可以做到这一点的ScriptModule。不幸的是，诸如控制流之类的东西被抹去了。

我们如何在TorchScript中忠实地表示此模块？我们提供了一个脚本编译器，它可以直接分析您的Python源代码以将其转换为TorchScript。让我们MyDecisionGate使用脚本编译器进行转换：

  1.     scripted_gate = torch.jit.script(MyDecisionGate())
  3.     my_cell = MyCell(scripted_gate)
  4.     traced_cell = torch.jit.script(my_cell)
  5.     print(traced_cell.code)

Out:

  1.     def forward(self,
  2.         x: Tensor,
  3.         h: Tensor) -> Tuple[Tensor, Tensor]:
  4.       _0 = self.linear
  5.       _1 = _0.weight
  6.       _2 = _0.bias
  7.       if torch.eq(torch.dim(x), 2):
  8.         _3 = torch.__isnot__(_2, None)
  9.       else:
  10.         _3 = False
  11.       if _3:
  12.         bias = ops.prim.unchecked_unwrap_optional(_2)
  13.         ret = torch.addmm(bias, x, torch.t(_1), beta=1, alpha=1)
  14.       else:
  15.         output = torch.matmul(x, torch.t(_1))
  16.         if torch.__isnot__(_2, None):
  17.           bias0 = ops.prim.unchecked_unwrap_optional(_2)
  18.           output0 = torch.add_(output, bias0, alpha=1)
  19.         else:
  20.           output0 = output
  21.         ret = output0
  22.       _4 = torch.gt(torch.sum(ret, dtype=None), 0)
  23.       if bool(_4):
  24.         _5 = ret
  25.       else:
  26.         _5 = torch.neg(ret)
  27.       new_h = torch.tanh(torch.add(_5, h, alpha=1))
  28.       return (new_h, new_h)

万岁！现在，我们已经忠实地捕获了我们在TorchScript中程序的行为。现在让我们尝试运行该程序：

  1.     # New inputs
  2.     x, h = torch.rand(3, 4), torch.rand(3, 4)
  3.     traced_cell(x, h)

混合脚本和跟踪

在某些情况下，我们只需要追踪的的结果而不需要全部脚本(例如，模块具有许多条件分支，这些分支我们并不希望展现在TorchScript中）。在这种情况下，脚本可以与用以下方法跟踪：torch.jit.script。他将只会追踪方法内的脚本，不会展示方法外的脚本情况。

第一种情况的一个示例：

  1.     class MyRNNLoop(torch.nn.Module):
  2.         def __init__(self):
  3.             super(MyRNNLoop, self).__init__()
  4.             self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))
  6.         def forward(self, xs):
  7.             h, y = torch.zeros(3, 4), torch.zeros(3, 4)
  8.             for i in range(xs.size(0)):
  9.                 y, h = self.cell(xs[i], h)
  10.             return y, h
  12.     rnn_loop = torch.jit.script(MyRNNLoop())
  13.     print(rnn_loop.code)

Out:

  1.     def forward(self,
  2.         xs: Tensor) -> Tuple[Tensor, Tensor]:
  3.       h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  4.       y = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  5.       y0, h0 = y, h
  6.       for i in range(torch.size(xs, 0)):
  7.         _0 = self.cell
  8.         _1 = torch.select(xs, 0, i)
  9.         _2 = _0.linear
  10.         weight = _2.weight
  11.         bias = _2.bias
  12.         _3 = torch.addmm(bias, _1, torch.t(weight), beta=1, alpha=1)
  13.         _4 = torch.gt(torch.sum(_3, dtype=None), 0)
  14.         if bool(_4):
  15.           _5 = _3
  16.         else:
  17.           _5 = torch.neg(_3)
  18.         _6 = torch.tanh(torch.add(_5, h0, alpha=1))
  19.         y0, h0 = _6, _6
  20.       return (y0, h0)

还有第二种情况的示例：

  1.     class WrapRNN(torch.nn.Module):
  2.       def __init__(self):
  3.         super(WrapRNN, self).__init__()
  4.         self.loop = torch.jit.script(MyRNNLoop())
  6.       def forward(self, xs):
  7.         y, h = self.loop(xs)
  8.         return torch.relu(y)
  10.     traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
  11.     print(traced.code)

Out:

  1.     def forward(self,
  2.         argument_1: Tensor) -> Tensor:
  3.       _0 = self.loop
  4.       h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  5.       h0 = h
  6.       for i in range(torch.size(argument_1, 0)):
  7.         _1 = _0.cell
  8.         _2 = torch.select(argument_1, 0, i)
  9.         _3 = _1.linear
  10.         weight = _3.weight
  11.         bias = _3.bias
  12.         _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
  13.         _5 = torch.gt(torch.sum(_4, dtype=None), 0)
  14.         if bool(_5):
  15.           _6 = _4
  16.         else:
  17.           _6 = torch.neg(_4)
  18.         h0 = torch.tanh(torch.add(_6, h0, alpha=1))
  19.       return torch.relu(h0)

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

保存和加载模型

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

  1.     traced.save('wrapped_rnn.zip')
  3.     loaded = torch.jit.load('wrapped_rnn.zip')
  5.     print(loaded)
  6.     print(loaded.code)

Out:

  1.     ScriptModule(
  2.       (loop): ScriptModule(
  3.         (cell): ScriptModule(
  4.           (dg): ScriptModule()
  5.           (linear): ScriptModule()
  6.         )
  7.       )
  8.     )
  9.     def forward(self,
  10.         argument_1: Tensor) -> Tensor:
  11.       _0 = self.loop
  12.       h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  13.       h0 = h
  14.       for i in range(torch.size(argument_1, 0)):
  15.         _1 = _0.cell
  16.         _2 = torch.select(argument_1, 0, i)
  17.         _3 = _1.linear
  18.         weight = _3.weight
  19.         bias = _3.bias
  20.         _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
  21.         _5 = torch.gt(torch.sum(_4, dtype=None), 0)
  22.         if bool(_5):
  23.           _6 = _4
  24.         else:
  25.           _6 = torch.neg(_4)
  26.         h0 = torch.tanh(torch.add(_6, h0, alpha=1))
  27.       return torch.relu(h0)

正如你所看到的，序列化保留了模块层次结构和我们一直在研究的代码。例如，也可以将模型加载到C ++中以实现不依赖Python的执行。

进一步阅读

我们已经完成了教程！有关更多涉及的演示，请查看NeurIPS演示，以使用TorchScript转换机器翻译模型

脚本的总运行时间： (0分钟0.252秒）