PyTorch 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 ++）中运行。

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

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

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

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

得出：

1.4.0

PyTorch 模型创作的基础

首先定义一个简单的Module。 Module是 PyTorch 中组成的基本单位。它包含了：

构造函数，为调用准备模块
一组Parameters和子Modules。这些由构造函数初始化，并且可以在调用期间由模块使用。
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.9541, 0.7233, 0.4907, 0.6169],
        [0.9117, 0.2329, 0.2512, 0.7751],
        [0.2949, 0.2434, 0.8694, 0.4242]]), tensor([[0.9541, 0.7233, 0.4907, 0.6169],
        [0.9117, 0.2329, 0.2512, 0.7751],
        [0.2949, 0.2434, 0.8694, 0.4242]]))

因此，我们已经：

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

我们实例化了该模块，并制作了x和y，它们只是 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.2940,  0.0822, -0.1697,  0.6644],
        [ 0.3065, -0.1165,  0.3684,  0.4877],
        [ 0.0409,  0.2764,  0.4881,  0.5211]], grad_fn=<TanhBackward>), tensor([[ 0.2940,  0.0822, -0.1697,  0.6644],
        [ 0.3065, -0.1165,  0.3684,  0.4877],
        [ 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 。简而言之，该系统允许我们通过潜在的复杂程序来计算导数。该设计为模型创作提供了极大的灵活性。

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

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.7407,  0.4486,  0.2651, -0.0298],
        [ 0.8582,  0.3146,  0.1919, -0.1760],
        [ 0.6428,  0.0017,  0.1307, -0.1543]], grad_fn=<TanhBackward>), tensor([[ 0.7407,  0.4486,  0.2651, -0.0298],
        [ 0.8582,  0.3146,  0.1919, -0.1760],
        [ 0.6428,  0.0017,  0.1307, -0.1543]], grad_fn=<TanhBackward>))

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

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

How autograd works

autograd 的工作原理

TorchScript 的基础

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

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

追踪`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)
)

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

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

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

print(traced_cell.graph)

得出:

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

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

print(traced_cell.code)

得出:

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

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

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

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

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

得出:

(tensor([[0.8188, 0.8444, 0.6618, 0.5024],
        [0.8359, 0.2624, 0.7421, 0.1236],
        [0.7331, 0.5259, 0.6288, 0.5338]], grad_fn=<TanhBackward>), tensor([[0.8188, 0.8444, 0.6618, 0.5024],
        [0.8359, 0.2624, 0.7421, 0.1236],
        [0.7331, 0.5259, 0.6288, 0.5338]], grad_fn=<TanhBackward>))
(tensor([[0.8188, 0.8444, 0.6618, 0.5024],
        [0.8359, 0.2624, 0.7421, 0.1236],
        [0.7331, 0.5259, 0.6288, 0.5338]],
       grad_fn=<DifferentiableGraphBackward>), tensor([[0.8188, 0.8444, 0.6618, 0.5024],
        [0.8359, 0.2624, 0.7421, 0.1236],
        [0.7331, 0.5259, 0.6288, 0.5338]],
       grad_fn=<DifferentiableGraphBackward>))

使用脚本转换模块

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

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.code)

得出:

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

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

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

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


my_cell = MyCell(scripted_gate)
traced_cell = torch.jit.script(my_cell)
print(traced_cell.code)

得出:

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

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

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

混合脚本和跟踪

在某些情况下，需要使用跟踪而不是脚本(例如，一个模块具有许多基于不变的 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], dtype=None, layout=None, device=None, pin_memory=None)
  y = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y0 = y
  h0 = h
  for i in range(torch.size(xs, 0)):
    _0 = (self.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,
    argument_1: Tensor) -> Tensor:
  _0, h, = (self.loop).forward(argument_1, )
  return torch.relu(h)

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

保存和加载模型

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

traced.save('wrapped_rnn.zip')


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


print(loaded)
print(loaded.code)

得出:

RecursiveScriptModule(
  original_name=Module
  (loop): RecursiveScriptModule(
    original_name=MyRNNLoop
    (cell): RecursiveScriptModule(
      original_name=Module
      (dg): RecursiveScriptModule(original_name=MyDecisionGate)
      (linear): RecursiveScriptModule(original_name=Module)
    )
  )
)
def forward(self,
    argument_1: Tensor) -> Tensor:
  _0, h, = (self.loop).forward(argument_1, )
  return torch.relu(h)

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

进一步阅读

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

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

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

w3cschool 编程狮，随时随地学编程

PyTorch TorchScript 简介

PyTorch 模型创作的基础

TorchScript 的基础

追踪`Modules`

使用脚本转换模块

混合脚本和跟踪

保存和加载模型

进一步阅读

PyTorch 入门

PyTorch深度学习

PyTorch 图片

Pytorch 音频

Pytorch 文本

PyTorch 命名为 Tensor(实验性）

PyTorch 强化学习

PyTorch 在生产中部署 PyTorch 模型

PyTorch 并行和分布式训练

PyTorch 扩展

PyTorch 模型优化

PyTorch 用其他语言

PyTorch 基础知识

PyTorch 笔记

PyTorch 语言绑定

Python API

PyTorch torchvision参考

PyTorch 音频参考

PyTorch 社区

w3cschool 编程狮，随时随地学编程

PyTorch TorchScript 简介

PyTorch 模型创作的基础

TorchScript 的基础

追踪Modules

使用脚本转换模块

混合脚本和跟踪

保存和加载模型

进一步阅读

PyTorch 入门

PyTorch深度学习

PyTorch 图片

Pytorch 音频

Pytorch 文本

PyTorch 命名为 Tensor(实验性）

PyTorch 强化学习

PyTorch 在生产中部署 PyTorch 模型

PyTorch 并行和分布式训练

PyTorch 扩展

PyTorch 模型优化

PyTorch 用其他语言

PyTorch 基础知识

PyTorch 笔记

PyTorch 语言绑定

Python API

PyTorch torchvision参考

PyTorch 音频参考

PyTorch 社区

追踪`Modules`