PyTorch TorchScript 簡介

2020-09-09 15:11 更新
原文： 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)作基礎(chǔ)，包括：
  2. 模組
  3. 定義forward功能
  4. 將模塊組成模塊的層次結(jié)構(gòu)
  5. 將 PyTorch 模塊轉(zhuǎn)換為 TorchScript(我們的高性能部署運行時）的特定方法
  6. 跟蹤現(xiàn)有模塊
  7. 使用腳本直接編譯模塊
  8. 如何組合這兩種方法
  9. 保存和加載 TorchScript 模塊

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

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

得出：

  1. 1.4.0

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

首先定義一個簡單的Module。 Module是 PyTorch 中組成的基本單位。 它包含了：

  1. 構(gòu)造函數(shù)，為調(diào)用準(zhǔn)備模塊
  2. 一組Parameters和子Modules。 這些由構(gòu)造函數(shù)初始化，并且可以在調(diào)用期間由模塊使用。
  3. forward功能。 這是調(diào)用模塊時運行的代碼。

我們來看一個小例子：

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

得出:

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

因此，我們已經(jīng)：

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

我們實例化了該模塊，并制作了x和y，它們只是 3x4 隨機值矩陣。 然后，我們使用my_cell(x, h)調(diào)用該單元格。 這依次調(diào)用我們的forward函數(shù)。

讓我們做一些更有趣的事情：

  1. class MyCell(torch.nn.Module):
  2.     def __init__(self):
  3.         super(MyCell, self).__init__()
  4.         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
  12. my_cell = MyCell()
  13. print(my_cell)
  14. print(my_cell(x, h))

得出：

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

我們已經(jīng)重新定義了模塊MyCell，但是這次我們添加了self.linear屬性，并在 forward 函數(shù)中調(diào)用了self.linear。

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

Module上的print將直觀地表示Module的子類層次結(jié)構(gòu)。 在我們的示例中，我們可以看到Linear子類及其參數(shù)。

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

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

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

  1. class MyDecisionGate(torch.nn.Module):
  2.     def forward(self, x):
  3.         if x.sum() > 0:
  4.             return x
  5.         else:
  6.             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)
  16.     def forward(self, x, h):
  17.         new_h = torch.tanh(self.dg(self.linear(x)) + h)
  18.         return new_h, new_h
  21. my_cell = MyCell()
  22. print(my_cell)
  23. print(my_cell(x, h))

得出:

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

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

How autograd works

autograd 的工作原理

TorchScript 的基礎(chǔ)

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

簡而言之，即使 PyTorch 具有靈活和動態(tài)的特性，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)
  7.     def forward(self, x, h):
  8.         new_h = torch.tanh(self.linear(x) + h)
  9.         return new_h, new_h
  12. my_cell = MyCell()
  13. x, h = torch.rand(3, 4), torch.rand(3, 4)
  14. traced_cell = torch.jit.trace(my_cell, (x, h))
  15. print(traced_cell)
  16. traced_cell(x, h)

得出:

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

我們倒退了一點，并學(xué)習(xí)了MyCell類的第二個版本。 和以前一樣，我們實例化了它，但是這一次，我們調(diào)用了torch.jit.trace，并傳入了Module，并傳入了示例輸入，網(wǎng)絡(luò)可能會看到。

這到底是做什么的？ 它調(diào)用了Module，記錄了運行Module時發(fā)生的操作，并創(chuàng)建了torch.jit.ScriptModule的實例(其中TracedModule是實例）

TorchScript 在中間表示(或 IR）中記錄其定義，在深度學(xué)習(xí)中通常將其稱為圖。 我們可以使用.graph屬性檢查圖形：

  1. print(traced_cell.graph)

得出:

  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)

但是，這是一個非常低級的表示形式，圖中包含的大多數(shù)信息對最終用戶沒有用。 相反，我們可以使用.code屬性給出代碼的 Python 語法解釋：

  1. print(traced_cell.code)

得出:

  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 代碼可以在其自己的解釋器中調(diào)用，該解釋器基本上是受限制的 Python 解釋器。 該解釋器不獲取全局解釋器鎖定，因此可以在同一實例上同時處理許多請求。
  2. 這種格式允許我們將整個模型保存到磁盤上，然后將其加載到另一個環(huán)境中，例如在以 Python 以外的語言編寫的服務(wù)器中
  3. TorchScript 為我們提供了一種表示形式，其中我們可以對代碼進行編譯器優(yōu)化以提供更有效的執(zhí)行
  4. TorchScript 允許我們與許多后端/設(shè)備運行時進行交互，與單個操作員相比，它們要求更廣泛的程序視圖。

我們可以看到，調(diào)用traced_cell會產(chǎn)生與 Python 模塊相同的結(jié)果：

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

得出:

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

使用腳本轉(zhuǎn)換模塊

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

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

得出:

  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輸出，我們可以發(fā)現(xiàn)找不到if-else分支！ 為什么？ 跟蹤完全按照我們所說的去做：運行代碼，記錄發(fā)生的操作，并構(gòu)造一個可以做到這一點的 ScriptModule。 不幸的是，諸如控制流之類的東西被抹去了。

我們?nèi)绾卧?TorchScript 中忠實地表示此模塊？ 我們提供了腳本編譯器，它可以直接分析您的 Python 源代碼以將其轉(zhuǎn)換為 TorchScript。 讓我們使用腳本編譯器轉(zhuǎn)換MyDecisionGate：

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

得出:

  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)

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

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

混合腳本和跟蹤

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

第一種情況的示例：

  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))
  7.     def forward(self, xs):
  8.         h, y = torch.zeros(3, 4), torch.zeros(3, 4)
  9.         for i in range(xs.size(0)):
  10.             y, h = self.cell(xs[i], h)
  11.         return y, h
  14. rnn_loop = torch.jit.script(MyRNNLoop())
  15. print(rnn_loop.code)

得出:

  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())
  7.     def forward(self, xs):
  8.         y, h = self.loop(xs)
  9.         return torch.relu(y)
  12. traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
  13. print(traced.code)

得出:

  1. def forward(self,
  2.     argument_1: Tensor) -> Tensor:
  3.   _0, h, = (self.loop).forward(argument_1, )
  4.   return torch.relu(h)

這樣，當(dāng)情況需要它們時，可以使用腳本和跟蹤并將它們一起使用。

保存和加載模型

我們提供 API，以存檔格式將 TorchScript 模塊保存到磁盤或從磁盤加載 TorchScript 模塊。 這種格式包括代碼，參數(shù)，屬性和調(diào)試信息，這意味著歸檔文件是模型的獨立表示，可以在完全獨立的過程中加載。 讓我們保存并加載包裝好的 RNN 模塊：

  1. traced.save('wrapped_rnn.zip')
  4. loaded = torch.jit.load('wrapped_rnn.zip')
  7. print(loaded)
  8. print(loaded.code)

得出:

  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)

如您所見，序列化保留了模塊層次結(jié)構(gòu)和我們一直在研究的代碼。 也可以將模型加載到中，例如到 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

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


以上內(nèi)容是否對您有幫助：
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號
微信公眾號

編程獅公眾號