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的子類）的中間表示形式，可以在高性能環(huán)境(例如 C ++）中運行。

在本教程中，我們將介紹：

1. PyTorch 中的模型創(chuàng)作基礎，包括：
2. 模組
3. 定義forward功能
4. 將模塊組成模塊的層次結構
5. 將 PyTorch 模塊轉換為 TorchScript(我們的高性能部署運行時）的特定方法
6. 跟蹤現(xiàn)有模塊
7. 使用腳本直接編譯模塊
8. 如何組合這兩種方法
9. 保存和加載 TorchScript 模塊

我們希望在完成本教程后，您將繼續(xù)學習和后續(xù)教程，該教程將引導您完成一個從 C ++實際調用 TorchScript 模型的示例。

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

得出：

1.4.0

PyTorch 模型創(chuàng)作的基礎

首先定義一個簡單的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.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]]))

因此，我們已經：

1. 創(chuàng)建了一個子類torch.nn.Module的類。
2. 定義一個構造函數。 構造函數沒有做太多事情，只是調用super的構造函數。
3. 定義了forward函數，該函數具有兩個輸入并返回兩個輸出。 forward函數的實際內容并不是很重要，但它是一種偽造的 RNN 單元格，即，該函數應用于循環(huán)。

我們實例化了該模塊，并制作了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。

這里到底發(fā)生了什么？ torch.nn.Linear是 PyTorch 標準庫中的Module。 就像MyCell一樣，可以使用調用語法來調用它。 我們正在建立Module的層次結構。

Module上的print將直觀地表示Module的子類層次結構。 在我們的示例中，我們可以看到Linear子類及其參數。

通過以這種方式組成Module，我們可以簡潔易讀地編寫具有可重用組件的模型。

您可能已經在輸出上注意到grad_fn。 這是 PyTorch 自動區(qū)分方法的詳細信息，稱為 autograd 。 簡而言之，該系統(tǒng)允許我們通過潛在的復雜程序來計算導數。 該設計為模型創(chuàng)作提供了極大的靈活性。

現(xiàn)在讓我們檢查一下靈活性：

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。 該模塊利用控制流。 控制流包括循環(huán)和if語句之類的內容。

給定完整的程序表示形式，許多框架都采用計算符號導數的方法。 但是，在 PyTorch 中，我們使用漸變色帶。 我們記錄發(fā)生的操作，并在計算派生時向后回放。 這樣，框架不必為語言中的所有構造顯式定義派生類。

autograd 的工作原理

TorchScript 的基礎

現(xiàn)在，讓我們以正在運行的示例為例，看看如何應用 TorchScript。

簡而言之，即使 PyTorch 具有靈活和動態(tài)的特性，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時發(fā)生的操作，并創(chuàng)建了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)

那么為什么我們要進行所有這些操作？ 有以下幾個原因：

1. TorchScript 代碼可以在其自己的解釋器中調用，該解釋器基本上是受限制的 Python 解釋器。 該解釋器不獲取全局解釋器鎖定，因此可以在同一實例上同時處理許多請求。
2. 這種格式允許我們將整個模型保存到磁盤上，然后將其加載到另一個環(huán)境中，例如在以 Python 以外的語言編寫的服務器中
3. TorchScript 為我們提供了一種表示形式，其中我們可以對代碼進行編譯器優(yōu)化以提供更有效的執(zhí)行
4. 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>))

使用腳本轉換模塊

原因是我們使用了模塊的第二版，而不是使用帶有控制流的子模塊的第二版。 現(xiàn)在讓我們檢查一下：

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輸出，我們可以發(fā)現(xiàn)找不到if-else分支！ 為什么？ 跟蹤完全按照我們所說的去做：運行代碼，記錄發(fā)生的操作，并構造一個可以做到這一點的 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)

萬歲！ 現(xiàn)在，我們已經忠實地捕獲了我們在 TorchScript 中程序的行為。 現(xiàn)在，讓我們嘗試運行該程序：

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

混合腳本和跟蹤

在某些情況下，需要使用跟蹤而不是腳本(例如，一個模塊具有許多基于不變的 Python 值做出的架構決策，而我們不希望它們出現(xiàn)在 TorchScript 中）。 在這種情況下，可以通過跟蹤來組成腳本：torch.jit.script將內聯(lián)被跟蹤模塊的代碼，而跟蹤將內聯(lián)腳本模塊的代碼。

第一種情況的示例：

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 ++ 中，以實現(xiàn)不依賴 Python 的執(zhí)行。

進一步閱讀

腳本的總運行時間：(0 分鐘 0.251 秒）

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

由獅身人面像畫廊生成的畫廊