TorchScript 简介

原文： https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html

注意

单击此处的下载完整的示例代码

James Reed(jamesreed@fb.com），Michael Suo(suo@fb.com），rev2

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

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

  1. PyTorch 中的模型创作基础，包括：
  • 模组
  • 定义forward功能
  • 将模块组成模块的层次结构
  1. 将 PyTorch 模块转换为 TorchScript(我们的高性能部署运行时）的特定方法
  • 跟踪现有模块
  • 使用脚本直接编译模块
  • 如何组合这两种方法
  • 保存和加载 TorchScript 模块

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

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

出：

  1. 1.4.0

PyTorch 模型创作的基础

首先定义一个简单的Module。 Module是 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.9541, 0.7233, 0.4907, 0.6169],
  2.         [0.9117, 0.2329, 0.2512, 0.7751],
  3.         [0.2949, 0.2434, 0.8694, 0.4242]]), tensor([[0.9541, 0.7233, 0.4907, 0.6169],
  4.         [0.9117, 0.2329, 0.2512, 0.7751],
  5.         [0.2949, 0.2434, 0.8694, 0.4242]]))

因此，我们已经：

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

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

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

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

Out:

  1. MyCell(
  2.   (linear): Linear(in_features=4, out_features=4, bias=True)
  3. )
  4. (tensor([[ 0.2940,  0.0822, -0.1697,  0.6644],
  5.         [ 0.3065, -0.1165,  0.3684,  0.4877],
  6.         [ 0.0409,  0.2764,  0.4881,  0.5211]], grad_fn=<TanhBackward>), tensor([[ 0.2940,  0.0822, -0.1697,  0.6644],
  7.         [ 0.3065, -0.1165,  0.3684,  0.4877],
  8.         [ 0.0409,  0.2764,  0.4881,  0.5211]], grad_fn=<TanhBackward>))

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

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

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

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

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

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

  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):
  10.         super(MyCell, self).__init__()
  11.         self.dg = MyDecisionGate()
  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()
  19. print(my_cell)
  20. print(my_cell(x, h))

Out:

  1. MyCell(
  2.   (dg): MyDecisionGate()
  3.   (linear): Linear(in_features=4, out_features=4, bias=True)
  4. )
  5. (tensor([[ 0.7407,  0.4486,  0.2651, -0.0298],
  6.         [ 0.8582,  0.3146,  0.1919, -0.1760],
  7.         [ 0.6428,  0.0017,  0.1307, -0.1543]], grad_fn=<TanhBackward>), tensor([[ 0.7407,  0.4486,  0.2651, -0.0298],
  8.         [ 0.8582,  0.3146,  0.1919, -0.1760],
  9.         [ 0.6428,  0.0017,  0.1307, -0.1543]], grad_fn=<TanhBackward>))

我们再次重新定义了 MyCell 类，但在这里我们定义了MyDecisionGate。 该模块利用控制流。 控制流包括循环和if语句之类的内容。

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

How autograd works

autograd 的工作原理

TorchScript 的基础

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

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

追踪Modules

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

Out:

  1. MyCell(
  2.   original_name=MyCell
  3.   (linear): Linear(original_name=Linear)
  4. )

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

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

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

  1. print(traced_cell.graph)

Out:

  1. graph(%self.1 : __torch__.torch.nn.modules.module.___torch_mangle_1.Module,
  2.       %input : Float(3, 4),
  3.       %h : Float(3, 4)):
  4.   %19 : __torch__.torch.nn.modules.module.Module = prim::GetAttr[name="linear"](%self.1)
  5.   %21 : Tensor = prim::CallMethod[name="forward"](%19, %input)
  6.   %12 : int = prim::Constant[value=1]() # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  7.   %13 : Float(3, 4) = aten::add(%21, %h, %12) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  8.   %14 : Float(3, 4) = aten::tanh(%13) # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  9.   %15 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%14, %14)
  10.   return (%15)

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

  1. print(traced_cell.code)

Out:

  1. def forward(self,
  2.     input: Tensor,
  3.     h: Tensor) -> Tuple[Tensor, Tensor]:
  4.   _0 = torch.add((self.linear).forward(input, ), h, alpha=1)
  5.   _1 = torch.tanh(_0)
  6.   return (_1, _1)

那么为什么我们要进行所有这些操作？ 有以下几个原因：

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

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

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

Out:

  1. (tensor([[0.8188, 0.8444, 0.6618, 0.5024],
  2.         [0.8359, 0.2624, 0.7421, 0.1236],
  3.         [0.7331, 0.5259, 0.6288, 0.5338]], grad_fn=<TanhBackward>), tensor([[0.8188, 0.8444, 0.6618, 0.5024],
  4.         [0.8359, 0.2624, 0.7421, 0.1236],
  5.         [0.7331, 0.5259, 0.6288, 0.5338]], grad_fn=<TanhBackward>))
  6. (tensor([[0.8188, 0.8444, 0.6618, 0.5024],
  7.         [0.8359, 0.2624, 0.7421, 0.1236],
  8.         [0.7331, 0.5259, 0.6288, 0.5338]],
  9.        grad_fn=<DifferentiableGraphBackward>), tensor([[0.8188, 0.8444, 0.6618, 0.5024],
  10.         [0.8359, 0.2624, 0.7421, 0.1236],
  11.         [0.7331, 0.5259, 0.6288, 0.5338]],
  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:

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

查看.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.dg).forward((self.linear).forward(x, ), )
  5.   new_h = torch.tanh(torch.add(_0, h, alpha=1))
  6.   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)

混合脚本和跟踪

在某些情况下，需要使用跟踪而不是脚本(例如，一个模块具有许多基于不变的 Python 值做出的架构决策，而我们不希望它们出现在 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 = y
  6.   h0 = h
  7.   for i in range(torch.size(xs, 0)):
  8.     _0 = (self.cell).forward(torch.select(xs, 0, i), h0, )
  9.     y1, h1, = _0
  10.     y0, h0 = y1, h1
  11.   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, h, = (self.loop).forward(argument_1, )
  4.   return torch.relu(h)

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

保存和加载模型

我们提供 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. RecursiveScriptModule(
  2.   original_name=Module
  3.   (loop): RecursiveScriptModule(
  4.     original_name=MyRNNLoop
  5.     (cell): RecursiveScriptModule(
  6.       original_name=Module
  7.       (dg): RecursiveScriptModule(original_name=MyDecisionGate)
  8.       (linear): RecursiveScriptModule(original_name=Module)
  9.     )
  10.   )
  11. )
  12. def forward(self,
  13.     argument_1: Tensor) -> Tensor:
  14.   _0, h, = (self.loop).forward(argument_1, )
  15.   return torch.relu(h)

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

进一步阅读

我们已经完成了教程！ 有关更多涉及的演示，请查看 NeurIPS 演示以使用 TorchScript 转换机器翻译模型： https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ

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

Download Python source code: Intro_to_TorchScript_tutorial.py Download Jupyter notebook: Intro_to_TorchScript_tutorial.ipynb

由狮身人面像画廊生成的画廊