TorchScript是什么？

TorchScript - PyTorch master documentation?pytorch.org

TorchScript是一種從PyTorch代碼創建可序列化和可優化模型的方法。任何TorchScript程序都可以從Python進程中保存，并加載到沒有Python依賴的進程中。

我們提供了一些工具來增量地將模型從純Python程序轉換為能夠獨立于Python運行的TorchScript程序，例如在獨立的c++程序中。這使得使用熟悉的Python工具在PyTorch中訓練模型，然后通過TorchScript將模型導出到生產環境中成為可能，在這種環境中，Python程序可能由于性能和多線程的原因不適用。

編寫TorchScript代碼

torch.jit.script(obj)

腳本化一個函數或者nn.Module對象，將會檢查它的源代碼， 將其作為TorchScript代碼使用TorchScrit編譯器編譯它，返回一個ScriptModule或ScriptFunction。 TorchScript語言自身是Python語言的一個子類, 因此它并非具有所有的Python語言特性。 torch.jit.script能夠被作為函數或裝飾器使用。參數obj可以是class, function, nn.Module。

具體地，腳本化一個函數： torch.jit.script 裝飾器將會通過編譯函數被裝飾函數體來構造一個ScriptFunction對象。例如：

import torch@torch.jit.script def foo(x, y):if x.max() > y.max():r = xelse:r = yreturn rprint(type(foo)) # torch.jit.ScriptFuncion# See the compiled graph as Python code print(foo.code)

腳本化一個nn.Module：默認地編譯其forward方法，并遞歸地編譯其子模塊以及被forward調用的函數。如果一個模塊只使用TorchScript中支持的特性，則不需要更改原始模塊代碼。編譯器將構建ScriptModule，其中包含原始模塊的屬性、參數和方法的副本。例如：

import torchclass MyModule(torch.nn.Module):def __init__(self, N, M):super(MyModule, self).__init__()# This parameter will be copied to the new ScriptModuleself.weight = torch.nn.Parameter(torch.rand(N, M))# When this submodule is used, it will be compiledself.linear = torch.nn.Linear(N, M)def forward(self, input):output = self.weight.mv(input)# This calls the `forward` method of the `nn.Linear` module, which will# cause the `self.linear` submodule to be compiled to a `ScriptModule` hereoutput = self.linear(output)return outputscripted_module = torch.jit.script(MyModule(2, 3))

編譯一個不在forward中的方法以及遞歸地編譯其內的所有方法，可在此方法上使用裝飾器torch.jit.export為了忽視某些方法也可以使用裝飾器為了忽視某些方法也可以使用裝飾器torch.jit.ignore和torch.jit.unused

import torch import torch.nn as nnclass MyModule(nn.Module):def __init__(self):super(MyModule, self).__init__()@torch.jit.exportdef some_entry_point(self, input):return input + 10@torch.jit.ignoredef python_only_fn(self, input):# This function won't be compiled, so any# Python APIs can be usedimport pdbpdb.set_trace()def forward(self, input):if self.training:self.python_only_fn(input)return input * 99scripted_module = torch.jit.script(MyModule()) print(scripted_module.some_entry_point(torch.randn(2, 2))) print(scripted_module(torch.randn(2, 2)))

torch.jit.trace(func,example_inputs,optimize=None,check_trace=True,check_inputs=None,check_tolerance=1e-5)

跟蹤一個函數并返回一個可執行的或ScriptFunction對象，將使用即時編譯(JIT)進行優化。跟蹤非常適合那些只操作單張量或張量的列表、字典和元組的代碼。使用torch.jit.trace和torch.jit.trace_module ，你能將一個模型或python函數轉為TorchScript中的ScriptModule或ScriptFunction。根據你提供的輸入樣例，它將會運行 該函數并記錄所有張量上執行的操作。

Tracing 僅僅正確地記錄那些不是數據依賴的函數和nn.Module（例如沒有對數據的條件判斷） 并且它們也沒有任何未跟蹤的外部依賴（例如執行輸入輸出或訪問全局變量). Tracing 只記錄在給定張量上運行給定函數時所執行的操作。 因此，返回的ScriptModule將始終在任何輸入上運行相同的跟蹤圖。當你的模塊需要根據輸入和/或模塊狀態運行不同的操作集時，這就產生了一些重要的影響。例如：

• Tracing不會記錄任何類似if語句或循環的控制流。當這個控制流在您的模塊中是常量時，這是沒有問題的，并且它通常內聯了控制流決策。但有時控制流實際上是模型本身的一部分。例如，一個遞歸網絡是一個輸入序列長度(可能是動態的)的循環。
• 在返回的ScriptModule中，無論ScriptModule處于哪種模式，在train和eval模式中具有不同行為的操作都將始終表現為處于跟蹤時所處的模式。

在這種情況下，Trace是不合適的，Script是更好的選擇。如果你跟蹤這樣的模型，您可能會在后續的模型調用中得到不正確的結果。當執行可能導致產生錯誤跟蹤的操作時，跟蹤程序將嘗試發出警告。

tracing a function:

import torchdef foo(x, y):return 2 * x + y# Run `foo` with the provided inputs and record the tensor operations traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))# `traced_foo` can now be run with the TorchScript interpreter or saved # and loaded in a Python-free environment

tracing a existing module

import torch import torch.nn as nnclass Net(nn.Module):def __init__(self):super(Net, self).__init__()self.conv = nn.Conv2d(1, 1, 3)def forward(self, x):return self.conv(x)n = Net() example_weight = torch.rand(1, 1, 3, 3) example_forward_input = torch.rand(1, 1, 3, 3)# Trace a specific method and construct `ScriptModule` with # a single `forward` method module = torch.jit.trace(n.forward, example_forward_input)# Trace a module (implicitly traces `forward`) and construct a # `ScriptModule` with a single `forward` method module = torch.jit.trace(n, example_forward_input)

torch.jit.trace_module(mod,inputs,optimize=None,check_trace=True,check_inputs=None,check_tolerance=1e-5)

跟蹤一個模塊并返回一個可執行的ScriptModule，該腳本模塊將使用即時編譯進行優化。當一個模塊被傳遞到torch.jit.trace，只運行和跟蹤forward方法。使用trace_module，您可以為要跟蹤的示例輸入指定一個方法名字典(參見下面的example_input參數)。

import torch import torch.nn as nnclass Net(nn.Module):def __init__(self):super(Net, self).__init__()self.conv = nn.Conv2d(1, 1, 3)def forward(self, x):return self.conv(x)def weighted_kernel_sum(self, weight):return weight * self.conv.weightn = Net() example_weight = torch.rand(1, 1, 3, 3) example_forward_input = torch.rand(1, 1, 3, 3)# Trace a specific method and construct `ScriptModule` with # a single `forward` method module = torch.jit.trace(n.forward, example_forward_input)# Trace a module (implicitly traces `forward`) and construct a # `ScriptModule` with a single `forward` method module = torch.jit.trace(n, example_forward_input)# Trace specific methods on a module (specified in `inputs`), constructs # a `ScriptModule` with `forward` and `weighted_kernel_sum` methods inputs = {'forward' : example_forward_input, 'weighted_kernel_sum' : example_weight} module = torch.jit.trace_module(n, inputs)

class torch.jit.ScriptModule

ScriptModule 封裝一個c++接口中的torch::jit::Module類， 有下列屬性及方法：

• code 返回forward方法的內部圖的打印表示(具有有效的Python語法)
• graph返回forward方法的內部圖的字符串表示形式
• inlined_graph返回forward方法的內部圖的字符串表示形式。此圖將被預處理為內聯所有函數和方法調用。
• save(f,_extra_files=ExtraFilesMap{})

class torch.jit.ScriptFunction 與上者類似

torch.jit.save(m,f,_extra_files=ExtraFilesMap{})

保存此模塊的脫機版本，以便在單獨的進程中使用。所保存的模塊序列化此模塊的所有方法、子模塊、參數和屬性。它可以使用torch::jit::load(文件名)加載到c++ API中，也可以使用torch.jit.load加載到Python API中。為了能夠保存模塊，它必須不調用任何本機Python函數。這意味著所有子模塊也必須是ScriptModule的子類。所有模塊，不管它們的設備是什么，總是在加載過程中加載到CPU上。這與torch.load()的語義不同，將來可能會改變。

import torch import ioclass MyModule(torch.nn.Module):def forward(self, x):return x + 10m = torch.jit.script(MyModule())# Save to file torch.jit.save(m, 'scriptmodule.pt') # This line is equivalent to the previous m.save("scriptmodule.pt")# Save to io.BytesIO buffer buffer = io.BytesIO() torch.jit.save(m, buffer)# Save with extra files extra_files = torch._C.ExtraFilesMap() extra_files['foo.txt'] = 'bar' torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)

torch.jit.load(f,map_location=None,_extra_files=ExtraFilesMap{})

加載先前用torch.jit.save保存的ScriptModule或ScriptFunction所有之前保存的模塊，無論它們的設備是什么，都首先加載到CPU上，然后移動到它們保存的設備上。如果失敗(例如，因為運行時系統沒有特定的設備)，就會引發異常。

import torch import iotorch.jit.load('scriptmodule.pt')# Load ScriptModule from io.BytesIO object with open('scriptmodule.pt', 'rb') as f:buffer = io.BytesIO(f.read())# Load all tensors to the original device torch.jit.load(buffer)# Load all tensors onto CPU, using a device buffer.seek(0) torch.jit.load(buffer, map_location=torch.device('cpu'))# Load all tensors onto CPU, using a string buffer.seek(0) torch.jit.load(buffer, map_location='cpu')# Load with extra files. extra_files = torch._C.ExtraFilesMap() extra_files['foo.txt'] = 'bar' torch.jit.load('scriptmodule.pt', _extra_files=extra_files) print(extra_files['foo.txt'])

torch.jit.ignore(drop=False, **kwargs)

這個裝飾器向編譯器表明，一個函數或方法應該被忽略，并保留為Python函數。這允許您在模型中保留尚未與TorchScript兼容的代碼。如果從TorchScript調用，被忽略的函數將把調用分派給Python解釋器。函數被忽略的模型不能導出。使用drop=True參數時可以，但會拋出異常。最好使用torch.jit.unused

import torch import torch.nn as nnclass MyModule(nn.Module):@torch.jit.ignoredef debugger(self, x):import pdbpdb.set_trace()def forward(self, x):x += 10# The compiler would normally try to compile `debugger`,# but since it is `@ignore`d, it will be left as a call# to Pythonself.debugger(x)return xm = torch.jit.script(MyModule())# Error! The call `debugger` cannot be saved since it calls into Python m.save("m.pt")

使用torch.jit.ignore(drop=True), 這一方法已被torch.jit.unused替代。

import torch import torch.nn as nnclass MyModule(nn.Module):@torch.jit.ignore(drop=True)def training_method(self, x):import pdbpdb.set_trace()def forward(self, x):if self.training:self.training_method(x)return xm = torch.jit.script(MyModule())# This is OK since `training_method` is not saved, the call is replaced # with a `raise`. m.save("m.pt")

torch.jit.unused(fn)

這個裝飾器向編譯器表明，應該忽略一個函數或方法，并用引發異常來替換它。這允許您在模型中保留與TorchScript不兼容的代碼，同時仍然導出模型。

import torch import torch.nn as nnclass MyModule(nn.Module):def __init__(self, use_memory_efficent):super(MyModule, self).__init__()self.use_memory_efficent = use_memory_efficent@torch.jit.unuseddef memory_efficient(self, x):import pdbpdb.set_trace()return x + 10def forward(self, x):# Use not-yet-scriptable memory efficient modeif self.use_memory_efficient:return self.memory_efficient(x)else:return x + 10m = torch.jit.script(MyModule(use_memory_efficent=False)) m.save("m.pt")m = torch.jit.script(MyModule(use_memory_efficient=True)) # exception raised m(torch.rand(100))

混合Tracing和Scripting

在許多情況下，跟蹤或腳本是將模型轉換為TorchScript的一種更簡單的方法。可以編寫跟蹤和腳本來滿足模型某一部分的特定需求。

腳本函數可以調用跟蹤函數。當您需要圍繞一個簡單的前饋模型使用控制流時，這一點特別有用。例如，序列到序列模型的波束搜索通常用腳本編寫，但可以調用使用跟蹤生成的編碼器模塊。

例如在腳本中調用跟蹤函數

import torchdef foo(x, y):return 2 * x + ytraced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))@torch.jit.script def bar(x):return traced_foo(x, x)

跟蹤函數也可以調用腳本函數。當模型的一小部分需要一些控制流時，這是很有用的，即使大部分模型只是一個前饋網絡。由跟蹤函數調用的腳本函數中的控制流被正確保存。

例如在跟蹤函數中調用腳本函數

import torch@torch.jit.script def foo(x, y):if x.max() > y.max():r = xelse:r = yreturn rdef bar(x, y, z):return foo(x, y) + ztraced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3)))

這個組合也適用于nn.Module。

import torch import torchvisionclass MyScriptModule(torch.nn.Module):def __init__(self):super(MyScriptModule, self).__init__()self.means = torch.nn.Parameter(torch.tensor([103.939, 116.779, 123.68]).resize_(1, 3, 1, 1))self.resnet = torch.jit.trace(torchvision.models.resnet18(),torch.rand(1, 3, 224, 224))def forward(self, input):return self.resnet(input - self.means)my_script_module = torch.jit.script(MyScriptModule())

總結

以上是生活随笔為你收集整理的pytorch forward_【Pytorch部署】TorchScript的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。