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__()
  4. def forward(self, x, h):
  5. new_h = torch.tanh(x + h)
  6. return new_h, new_h
  7. my_cell = MyCell()
  8. x = torch.rand(3, 4)
  9. h = torch.rand(3, 4)
  10. 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)。

我們實例化了該模塊,并制作了xy,它們只是 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)
  5. def forward(self, x, h):
  6. new_h = torch.tanh(self.linear(x) + h)
  7. return new_h, new_h
  8. my_cell = MyCell()
  9. print(my_cell)
  10. 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
  7. class MyCell(torch.nn.Module):
  8. def __init__(self):
  9. super(MyCell, self).__init__()
  10. self.dg = MyDecisionGate()
  11. self.linear = torch.nn.Linear(4, 4)
  12. def forward(self, x, h):
  13. new_h = torch.tanh(self.dg(self.linear(x)) + h)
  14. return new_h, new_h
  15. my_cell = MyCell()
  16. print(my_cell)
  17. 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)
  5. def forward(self, x, h):
  6. new_h = torch.tanh(self.linear(x) + h)
  7. return new_h, new_h
  8. my_cell = MyCell()
  9. x, h = torch.rand(3, 4), torch.rand(3, 4)
  10. traced_cell = torch.jit.trace(my_cell, (x, h))
  11. print(traced_cell)
  12. 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
  7. class MyCell(torch.nn.Module):
  8. def __init__(self, dg):
  9. super(MyCell, self).__init__()
  10. self.dg = dg
  11. self.linear = torch.nn.Linear(4, 4)
  12. def forward(self, x, h):
  13. new_h = torch.tanh(self.dg(self.linear(x)) + h)
  14. return new_h, new_h
  15. my_cell = MyCell(MyDecisionGate())
  16. traced_cell = torch.jit.trace(my_cell, (x, h))
  17. 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())
  2. my_cell = MyCell(scripted_gate)
  3. traced_cell = torch.jit.script(my_cell)
  4. 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))
  5. def forward(self, xs):
  6. h, y = torch.zeros(3, 4), torch.zeros(3, 4)
  7. for i in range(xs.size(0)):
  8. y, h = self.cell(xs[i], h)
  9. return y, h
  10. rnn_loop = torch.jit.script(MyRNNLoop())
  11. 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())
  5. def forward(self, xs):
  6. y, h = self.loop(xs)
  7. return torch.relu(y)
  8. traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
  9. 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')
  2. loaded = torch.jit.load('wrapped_rnn.zip')
  3. print(loaded)
  4. 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

公眾號
微信公眾號

編程獅公眾號