本教程是對TorchScript的簡介，TorchScript是PyTorch模型（nn.Module的子類）的中間表示，可以在高性能環境（例如C ++）中運行。

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

1. PyTorch中的模型創作基礎，包括： 模組 定義前向功能 將模塊組成模塊的層次結構
2. 將PyTorch模塊轉換為TorchScript（我們的高性能部署運行時）的特定方法 跟蹤現有模塊 使用腳本直接編譯模塊 如何組合這兩種方法 保存和加載TorchScript模塊

我們希望在完成本教程之後，您將繼續閱讀後續教程，該教程將引導您真正地從C ++調用TorchScript模型的示例。

import torch  # 這是同時使用PyTorch和TorchScript所需的全部導入！print(torch.__version__)

• 輸出結果

1.3.0

1.PyTorch模型基礎

讓我們開始定義一個簡單的模塊。模塊是PyTorch中組成的基本單位。它包含：

• 構造函數，為調用準備模塊
• 一組參數和子模塊。這些由構造函數初始化，並且可以在調用期間由模塊使用。
• 前進功能。這是調用模塊時運行的代碼。
  我們來看一個小例子：

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_hmy_cell = MyCell()x = torch.rand(3, 4)h = torch.rand(3, 4)print(my_cell(x, h))

• 輸出結果

(tensor([[0.5139, 0.6451, 0.3697, 0.7738],        [0.7936, 0.5864, 0.8063, 0.9324],        [0.6479, 0.8408, 0.8062, 0.7263]]), tensor([[0.5139, 0.6451, 0.3697, 0.7738],        [0.7936, 0.5864, 0.8063, 0.9324],        [0.6479, 0.8408, 0.8062, 0.7263]]))

因此，我們已經：

1. 創建了一個子類torch.nn.Module的類。
2. 定義一個構造函數。構造函數沒有做太多事情，只是將構造函數稱為super。
3. 定義了一個正向功能，該功能需要兩個輸入並返回兩個輸出。前向函數的實際內容並不是很重要，但是它是一種偽造的https://colah.github.io/posts/2015-08-Understanding-LSTMs/
   -即，該函數應用於循環。

我們實例化了該模塊，並製作了x和y，它們只是3x4的隨機值矩陣。 然後，我們使用my_cell（x，h）調用該單元格。 這又調用了我們的轉發功能。

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

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_hmy_cell = MyCell()print(my_cell)print(my_cell(x, h))

• 輸出結果

MyCell(  (linear): Linear(in_features=4, out_features=4, bias=True))(tensor([[ 0.3941,  0.4160, -0.1086,  0.8432],        [ 0.5604,  0.4003,  0.5009,  0.6842],        [ 0.7084,  0.7147,  0.1818,  0.8296]], grad_fn=), tensor([[ 0.3941,  0.4160, -0.1086,  0.8432],        [ 0.5604,  0.4003,  0.5009,  0.6842],        [ 0.7084,  0.7147,  0.1818,  0.8296]], grad_fn=))

我們已經重新定義了模塊MyCell，但是這次我們添加了self.linear屬性，並在前進（forward）函數中調用了self.linear。

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

在模塊上列印可以直觀地表示該模塊的子類層次結構。在我們的示例中，我們可以看到我們的線性子類及其參數。

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

您可能已經在輸出中注意到grad_fn。這是PyTorch的自動區分方法（稱為https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html）的詳細信息。簡而言之，該系統允許我們通過潛在的複雜程序來計算導數。該設計為模型創作提供了極大的靈活性。

現在，讓我們檢查一下它的靈活性：

class MyDecisionGate(torch.nn.Module):    def forward(self, x):        if x.sum() > 0:            return x        else:            return -xclass 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_hmy_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.0850, 0.2812, 0.5188, 0.8523],        [0.1233, 0.3948, 0.6615, 0.7466],        [0.7072, 0.6103, 0.6953, 0.7047]], grad_fn=), tensor([[0.0850, 0.2812, 0.5188, 0.8523],        [0.1233, 0.3948, 0.6615, 0.7466],        [0.7072, 0.6103, 0.6953, 0.7047]], grad_fn=))

我們再次重新定義了MyCell類，但是在這裡我們定義了MyDecisionGate。該模塊利用控制流程。控制流包括循環和if語句之類的東西。

給定完整的程序表示形式，許多框架都採用計算符號派生的方法。但是，在PyTorch中，我們使用梯度帶。我們記錄操作發生時的操作，並在計算衍生產品時向後回放。這樣，框架不必為語言中的所有構造顯式定義派生類。

2.TorchScript的基礎

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

簡而言之，即使PyTorch具有靈活和動態的特性，TorchScript也提供了捕獲模型定義的工具。讓我們開始研究所謂的跟蹤。

2.1 跟蹤（Tracing）模塊

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

• 輸出結果

TracedModule[MyCell](  original_name=MyCell  (linear): TracedModule[Linear](original_name=Linear))

我們倒退了一點，並選擇了MyCell類的第二個版本。和以前一樣，我們實例化了它，但是這次，我們調用了torch.jit.trace，在Module(模塊)中傳遞了該示例，並在示例中傳遞了網絡可能看到的輸入。

這到底是做什麼的？ 它已調用模塊，記錄了模塊運行時發生的操作，並創建了torch.jit.ScriptModule的實例（TracedModule是其實例）

TorchScript將其定義記錄在中間表示（或IR）中，在深度學習中通常稱為圖形。我們可以檢查帶有.graph屬性的圖：

print(traced_cell.graph)

• 輸出結果

graph(%self : ClassType,      %input : Float(3, 4),      %h : Float(3, 4)):  %1 : ClassType = prim::GetAttr[name="linear"](%self)  %weight : Tensor = prim::GetAttr[name="weight"](%1)  %bias : Tensor = prim::GetAttr[name="bias"](%1)  %6 : Float(4, 4) = aten::t(%weight), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0  %7 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0  %8 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0  %9 : Float(3, 4) = aten::addmm(%bias, %input, %6, %7, %8), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0  %10 : int = prim::Constant[value=1](), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0  %11 : Float(3, 4) = aten::add(%9, %h, %10), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0  %12 : Float(3, 4) = aten::tanh(%11), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0  %13 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%12, %12)  return (%13)

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

print(traced_cell.code)

• 輸出結果

import __torch__import __torch__.torch.nn.modules.lineardef forward(self,    input: Tensor,    h: Tensor) -> Tuple[Tensor, Tensor]:  _0 = self.linear  weight = _0.weight  bias = _0.bias  _1 = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)  _2 = torch.tanh(torch.add(_1, h, alpha=1))  return (_2, _2)

那麼為什麼我們要做所有這些呢？ 有以下幾個原因：

1. TorchScript代碼可以在其自己的解釋器中調用，該解釋器基本上是受限制的Python解釋器。該解釋器不被全局解釋器鎖定，因此可以在同
   一實例上同時處理許多請求。
2. 這種格式使我們可以將整個模型保存到磁碟上，並將其加載到另一個環境中，例如在以Python以外的語言編寫的伺服器中
3. TorchScript為我們提供了一種表示形式，其中我們可以對代碼進行編譯器優化以提供更有效的執行
4. TorchScript允許我們與許多後端/設備運行時進行接口，這些運行時比單個操作員需要更廣泛的程序視圖。

我們可以看到，調用traced_cell產生的結果與Python模塊相同：

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

• 輸出結果

(tensor([[-0.3983,  0.5954,  0.2587, -0.3748],        [-0.5033,  0.4471,  0.8264,  0.2135],        [ 0.3430,  0.5561,  0.6794, -0.2273]], grad_fn=), tensor([[-0.3983,  0.5954,  0.2587, -0.3748],        [-0.5033,  0.4471,  0.8264,  0.2135],        [ 0.3430,  0.5561,  0.6794, -0.2273]], grad_fn=))(tensor([[-0.3983,  0.5954,  0.2587, -0.3748],        [-0.5033,  0.4471,  0.8264,  0.2135],        [ 0.3430,  0.5561,  0.6794, -0.2273]],       grad_fn=), tensor([[-0.3983,  0.5954,  0.2587, -0.3748],        [-0.5033,  0.4471,  0.8264,  0.2135],        [ 0.3430,  0.5561,  0.6794, -0.2273]],       grad_fn=))

3.使用腳本轉換模塊

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

class MyDecisionGate(torch.nn.Module):    def forward(self, x):        if x.sum() > 0:            return x        else:            return -xclass 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_hmy_cell = MyCell(MyDecisionGate())traced_cell = torch.jit.trace(my_cell, (x, h))print(traced_cell.code)

• 輸出結果

import __torch__.___torch_mangle_0import __torch__import __torch__.torch.nn.modules.linear.___torch_mangle_1def forward(self,    input: Tensor,    h: Tensor) -> Tuple[Tensor, Tensor]:  _0 = self.linear  weight = _0.weight  bias = _0.bias  x = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)  _1 = torch.tanh(torch.add(x, h, alpha=1))  return (_1, _1)

查看.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)

• 輸出結果

import __torch__.___torch_mangle_3import __torch__.___torch_mangle_2import __torch__.torch.nn.modules.linear.___torch_mangle_4def forward(self,    x: Tensor,    h: Tensor) -> Tuple[Tensor, Tensor]:  _0 = self.linear  _1 = _0.weight  _2 = _0.bias  if torch.eq(torch.dim(x), 2):    _3 = torch.__isnot__(_2, None)  else:    _3 = False  if _3:    bias = ops.prim.unchecked_unwrap_optional(_2)    ret = torch.addmm(bias, x, torch.t(_1), beta=1, alpha=1)  else:    output = torch.matmul(x, torch.t(_1))    if torch.__isnot__(_2, None):      bias0 = ops.prim.unchecked_unwrap_optional(_2)      output0 = torch.add_(output, bias0, alpha=1)    else:      output0 = output    ret = output0  _4 = torch.gt(torch.sum(ret, dtype=None), 0)  if bool(_4):    _5 = ret  else:    _5 = torch.neg(ret)  new_h = torch.tanh(torch.add(_5, h, alpha=1))  return (new_h, new_h)

現在，我們已經忠實地捕獲了我們在TorchScript中程序的行為。 現在，讓我們嘗試運行該程序：

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

3.1 混合腳本(Scripting)和跟蹤(Tracing)

在某些情況下，需要使用跟蹤而不是腳本（例如，模塊具有許多架構決策，這些決策是基於我們希望不會出現在TorchScript中的恆定Python值做出的）。在這種情況下，可以通過跟蹤來編寫腳本：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, hrnn_loop = torch.jit.script(MyRNNLoop())print(rnn_loop.code)

• 輸出結果

import __torch__import __torch__.___torch_mangle_5import __torch__.___torch_mangle_2import __torch__.torch.nn.modules.linear.___torch_mangle_6def 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    _1 = torch.select(xs, 0, i)    _2 = _0.linear    weight = _2.weight    bias = _2.bias    _3 = torch.addmm(bias, _1, torch.t(weight), beta=1, alpha=1)    _4 = torch.gt(torch.sum(_3, dtype=None), 0)    if bool(_4):      _5 = _3    else:      _5 = torch.neg(_3)    _6 = torch.tanh(torch.add(_5, h0, alpha=1))    y0, h0 = _6, _6  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)

• 輸出結果

import __torch__import __torch__.___torch_mangle_9import __torch__.___torch_mangle_7import __torch__.___torch_mangle_2import __torch__.torch.nn.modules.linear.___torch_mangle_8def forward(self,    argument_1: Tensor) -> Tensor:  _0 = self.loop  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)  h0 = h  for i in range(torch.size(argument_1, 0)):    _1 = _0.cell    _2 = torch.select(argument_1, 0, i)    _3 = _1.linear    weight = _3.weight    bias = _3.bias    _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)    _5 = torch.gt(torch.sum(_4, dtype=None), 0)    if bool(_5):      _6 = _4    else:      _6 = torch.neg(_4)    h0 = torch.tanh(torch.add(_6, h0, alpha=1))  return torch.relu(h0)

這樣，當情況需要它們時，可以使用腳本和跟蹤並將它們一起使用。

4.保存和加載模型

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

traced.save('wrapped_rnn.zip')loaded = torch.jit.load('wrapped_rnn.zip')print(loaded)print(loaded.code)

• 輸出結果

ScriptModule(  original_name=WrapRNN  (loop): ScriptModule(    original_name=MyRNNLoop    (cell): ScriptModule(      original_name=MyCell      (dg): ScriptModule(original_name=MyDecisionGate)      (linear): ScriptModule(original_name=Linear)    )  ))import __torch__import __torch__.___torch_mangle_9import __torch__.___torch_mangle_7import __torch__.___torch_mangle_2import __torch__.torch.nn.modules.linear.___torch_mangle_8def forward(self,    argument_1: Tensor) -> Tensor:  _0 = self.loop  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)  h0 = h  for i in range(torch.size(argument_1, 0)):    _1 = _0.cell    _2 = torch.select(argument_1, 0, i)    _3 = _1.linear    weight = _3.weight    bias = _3.bias    _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)    _5 = torch.gt(torch.sum(_4, dtype=None), 0)    if bool(_5):      _6 = _4    else:      _6 = torch.neg(_4)    h0 = torch.tanh(torch.add(_6, h0, alpha=1))  return torch.relu(h0)

如您所見，序列化保留了模塊層次結構和我們一直在研究的代碼。例如，也可以將模型加載到C ++中以實現不依賴Python的執行。

進一步閱讀

我們已經完成了教程！有關更多涉及的演示，請查看NeurIPS演示，以使用TorchScript轉換機器翻譯模型：https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ

腳本的總運行時間：（0分鐘0.247秒）