TorchScript 語言參考¶
本參考手冊描述了 TorchScript 語言的語法和核心語義。TorchScript 是 Python 語言的一個靜態型別子集。本文件解釋了 TorchScript 中支援的 Python 特性，以及該語言與普通 Python 的區別。本參考手冊中未提及的任何 Python 特性均不屬於 TorchScript。TorchScript 特別關注在 PyTorch 中表示神經網路模型所需的 Python 特性。
術語¶
本文件使用以下術語
模式 |
註釋 |
|---|---|
|
表示給定符號的定義方式。 |
|
表示語法中的實際關鍵字和分隔符。 |
|
表示 A 或 B。 |
|
表示分組。 |
|
表示可選。 |
|
表示正則表示式，其中項 A 重複至少一次。 |
|
表示正則表示式，其中項 A 重複零次或多次。 |
型別系統¶
TorchScript 是 Python 的一個靜態型別子集。TorchScript 與完整 Python 語言的最大區別在於，TorchScript 只支援表達神經網路模型所需的一小部分型別。
TorchScript 型別¶
TorchScript 型別系統包含 TSType 和 TSModuleType，定義如下。
TSAllType ::= TSType | TSModuleType
TSType ::= TSMetaType | TSPrimitiveType | TSStructuralType | TSNominalType
TSType 表示大多數 TorchScript 型別，它們是可組合的，並且可以在 TorchScript 型別標註中使用。TSType 指代以下任何一種型別
元型別，例如
Any基本型別，例如
int,float和str結構型別，例如
Optional[int]或List[MyClass]名義型別（Python 類），例如
MyClass(使用者定義),torch.tensor(內建)
TSModuleType 表示 torch.nn.Module 及其子類。它與 TSType 的處理方式不同，因為其型別模式部分從物件例項推斷，部分從類定義推斷。因此，TSModuleType 的例項可能不遵循相同的靜態型別模式。出於型別安全考慮，TSModuleType 不能用作 TorchScript 型別標註，也不能與 TSType 組合使用。
元型別¶
元型別非常抽象，它們更像是型別約束而非具體型別。目前 TorchScript 定義了一種元型別 Any，它表示任何 TorchScript 型別。
Any 型別¶
Any 型別表示任何 TorchScript 型別。Any 不指定任何型別約束，因此對 Any 不進行型別檢查。因此，它可以繫結到任何 Python 或 TorchScript 資料型別（例如 int、TorchScript tuple 或任何未被 script 的 Python 類）。
TSMetaType ::= "Any"
其中
Any是 typing 模組中的 Python 類名。因此，要使用Any型別，必須從typing匯入它（例如from typing import Any）。由於
Any可以表示任何 TorchScript 型別，因此允許對該型別的值執行的運算子集是有限的。
Any 型別支援的運算子¶
對
Any型別資料的賦值。繫結到
Any型別的引數或返回值。x is,x is not，其中x是Any型別。isinstance(x, Type)，其中x是Any型別。Any型別的資料是可列印的。如果資料是相同型別
T的值列表且T支援比較運算子，則List[Any]型別的資料可能是可排序的。
與 Python 相比
Any 是 TorchScript 型別系統中約束最少的型別。從這個意義上說，它與 Python 中的 Object 類非常相似。但是，Any 只支援 Object 支援的運算子和方法的一個子集。
設計說明¶
當我們 script 一個 PyTorch 模組時，可能會遇到不參與指令碼執行的資料。儘管如此，它仍然必須由型別模式描述。描述未使用資料（在指令碼上下文中）的靜態型別不僅麻煩，還可能導致不必要的 scripting 失敗。Any 的引入是為了描述那些對於編譯而言不需要精確靜態型別的資料的型別。
示例 1
此示例說明了如何使用 Any 允許元組引數的第二個元素可以是任何型別。這是可能的，因為 x[1] 不涉及任何需要知道其精確型別的計算。
import torch
from typing import Tuple
from typing import Any
@torch.jit.export
def inc_first_element(x: Tuple[int, Any]):
return (x[0]+1, x[1])
m = torch.jit.script(inc_first_element)
print(m((1,2.0)))
print(m((1,(100,200))))
上面的示例產生以下輸出
(2, 2.0)
(2, (100, 200))
元組的第二個元素是 Any 型別，因此可以繫結到多種型別。例如，(1, 2.0) 將 float 型別繫結到 Any，就像在 Tuple[int, Any] 中一樣；而 (1, (100, 200)) 在第二次呼叫中將 tuple 繫結到 Any。
示例 2
此示例說明了如何使用 isinstance 動態檢查標註為 Any 型別的資料的型別
import torch
from typing import Any
def f(a:Any):
print(a)
return (isinstance(a, torch.Tensor))
ones = torch.ones([2])
m = torch.jit.script(f)
print(m(ones))
上面的示例產生以下輸出
1
1
[ CPUFloatType{2} ]
True
基本型別¶
TorchScript 基本型別是表示單一型別值並具有單一預定義型別名稱的型別。
TSPrimitiveType ::= "int" | "float" | "double" | "complex" | "bool" | "str" | "None"
結構型別¶
結構型別是結構化定義且沒有使用者定義名稱（與名義型別不同）的型別，例如 Future[int]。結構型別可以與任何 TSType 組合使用。
TSStructuralType ::= TSTuple | TSNamedTuple | TSList | TSDict |
TSOptional | TSUnion | TSFuture | TSRRef | TSAwait
TSTuple ::= "Tuple" "[" (TSType ",")* TSType "]"
TSNamedTuple ::= "namedtuple" "(" (TSType ",")* TSType ")"
TSList ::= "List" "[" TSType "]"
TSOptional ::= "Optional" "[" TSType "]"
TSUnion ::= "Union" "[" (TSType ",")* TSType "]"
TSFuture ::= "Future" "[" TSType "]"
TSRRef ::= "RRef" "[" TSType "]"
TSAwait ::= "Await" "[" TSType "]"
TSDict ::= "Dict" "[" KeyType "," TSType "]"
KeyType ::= "str" | "int" | "float" | "bool" | TensorType | "Any"
其中
Tuple,List,Optional,Union,Future,Dict表示在typing模組中定義的 Python 型別類名。要使用這些型別名，必須從typing匯入它們（例如from typing import Tuple）。namedtuple表示 Python 類collections.namedtuple或typing.NamedTuple。Future和RRef表示 Python 類torch.futures和torch.distributed.rpc。Await表示 Python 類torch._awaits._Await
與 Python 相比
除了可以與 TorchScript 型別組合外，這些 TorchScript 結構型別通常支援其 Python 對應型別的一組通用運算子和方法。
示例 1
此示例使用 typing.NamedTuple 語法定義元組
import torch
from typing import NamedTuple
from typing import Tuple
class MyTuple(NamedTuple):
first: int
second: int
def inc(x: MyTuple) -> Tuple[int, int]:
return (x.first+1, x.second+1)
t = MyTuple(first=1, second=2)
scripted_inc = torch.jit.script(inc)
print("TorchScript:", scripted_inc(t))
上面的示例產生以下輸出
TorchScript: (2, 3)
示例 2
此示例使用 collections.namedtuple 語法定義元組
import torch
from typing import NamedTuple
from typing import Tuple
from collections import namedtuple
_AnnotatedNamedTuple = NamedTuple('_NamedTupleAnnotated', [('first', int), ('second', int)])
_UnannotatedNamedTuple = namedtuple('_NamedTupleAnnotated', ['first', 'second'])
def inc(x: _AnnotatedNamedTuple) -> Tuple[int, int]:
return (x.first+1, x.second+1)
m = torch.jit.script(inc)
print(inc(_UnannotatedNamedTuple(1,2)))
上面的示例產生以下輸出
(2, 3)
示例 3
此示例說明了標註結構型別時的一個常見錯誤，即沒有從 typing 模組匯入複合型別類
import torch
# ERROR: Tuple not recognized because not imported from typing
@torch.jit.export
def inc(x: Tuple[int, int]):
return (x[0]+1, x[1]+1)
m = torch.jit.script(inc)
print(m((1,2)))
執行上述程式碼會產生以下 scripting 錯誤
File "test-tuple.py", line 5, in <module>
def inc(x: Tuple[int, int]):
NameError: name 'Tuple' is not defined
解決方法是在程式碼開頭新增 from typing import Tuple 這一行。
名義型別¶
TorchScript 名義型別是 Python 類。這些型別之所以稱為名義型別，是因為它們是使用自定義名稱宣告的，並使用類名進行比較。名義類進一步分為以下幾類
TSNominalType ::= TSBuiltinClasses | TSCustomClass | TSEnum
其中，TSCustomClass 和 TSEnum 必須可以編譯為 TorchScript 中間表示 (IR)。這是由型別檢查器強制執行的。
內建類¶
內建名義型別是其語義已內建到 TorchScript 系統中的 Python 類（例如 tensor 型別）。TorchScript 定義了這些內建名義型別的語義，並且通常只支援其 Python 類定義的方法或屬性的一個子集。
TSBuiltinClass ::= TSTensor | "torch.device" | "torch.Stream" | "torch.dtype" |
"torch.nn.ModuleList" | "torch.nn.ModuleDict" | ...
TSTensor ::= "torch.Tensor" | "common.SubTensor" | "common.SubWithTorchFunction" |
"torch.nn.parameter.Parameter" | and subclasses of torch.Tensor
關於 torch.nn.ModuleList 和 torch.nn.ModuleDict 的特別說明¶
儘管 torch.nn.ModuleList 和 torch.nn.ModuleDict 在 Python 中定義為列表和字典，但在 TorchScript 中它們更像元組
在 TorchScript 中，
torch.nn.ModuleList或torch.nn.ModuleDict的例項是不可變的。迭代
torch.nn.ModuleList或torch.nn.ModuleDict的程式碼會完全展開，以便torch.nn.ModuleList的元素或torch.nn.ModuleDict的鍵可以是torch.nn.Module的不同子類。
示例
以下示例重點介紹了幾個內建 Torchscript 類（torch.*）的使用
import torch
@torch.jit.script
class A:
def __init__(self):
self.x = torch.rand(3)
def f(self, y: torch.device):
return self.x.to(device=y)
def g():
a = A()
return a.f(torch.device("cpu"))
script_g = torch.jit.script(g)
print(script_g.graph)
自定義類¶
與內建類不同，自定義類的語義是使用者定義的，並且整個類定義必須可編譯為 TorchScript IR 並遵守 TorchScript 型別檢查規則。
TSClassDef ::= [ "@torch.jit.script" ]
"class" ClassName [ "(object)" ] ":"
MethodDefinition |
[ "@torch.jit.ignore" ] | [ "@torch.jit.unused" ]
MethodDefinition
其中
類必須是新式類。Python 3 只支援新式類。在 Python 2.x 中，透過繼承 object 來指定新式類。
例項資料屬性是靜態型別的，並且例項屬性必須在
__init__()方法內部透過賦值宣告。不支援方法過載（即，不能有多個同名方法）。
MethodDefinition必須可編譯為 TorchScript IR 並遵守 TorchScript 的型別檢查規則（即，所有方法都必須是有效的 TorchScript 函式，並且類屬性定義必須是有效的 TorchScript 語句）。torch.jit.ignore和torch.jit.unused可用於忽略未完全 torchscript 化或應被編譯器忽略的方法或函式。
與 Python 相比
與 Python 自定義類相比，TorchScript 自定義類有相當多的限制。TorchScript 自定義類
不支援類屬性。
除了繼承介面型別或 object 外，不支援子類化。
不支援方法過載。
必須在
__init__()中初始化其所有例項屬性；這是因為 TorchScript 透過在__init__()中推斷屬性型別來構建類的靜態模式。必須只包含滿足 TorchScript 型別檢查規則且可編譯為 TorchScript IRs 的方法。
示例 1
如果 Python 類用 @torch.jit.script 標註，就可以在 TorchScript 中使用，類似於 TorchScript 函式的宣告方式
@torch.jit.script
class MyClass:
def __init__(self, x: int):
self.x = x
def inc(self, val: int):
self.x += val
示例 2
TorchScript 自定義類型別必須透過在 __init__() 中賦值來“宣告”其所有例項屬性。如果一個例項屬性未在 __init__() 中定義，但在類的其他方法中訪問，則該類無法編譯為 TorchScript 類，如下例所示
import torch
@torch.jit.script
class foo:
def __init__(self):
self.y = 1
# ERROR: self.x is not defined in __init__
def assign_x(self):
self.x = torch.rand(2, 3)
該類將編譯失敗併發出以下錯誤
RuntimeError:
Tried to set nonexistent attribute: x. Did you forget to initialize it in __init__()?:
def assign_x(self):
self.x = torch.rand(2, 3)
~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
示例 3
在此示例中，TorchScript 自定義類定義了一個類變數名，這是不允許的
import torch
@torch.jit.script
class MyClass(object):
name = "MyClass"
def __init__(self, x: int):
self.x = x
def fn(a: MyClass):
return a.name
這會導致以下編譯時錯誤
RuntimeError:
'__torch__.MyClass' object has no attribute or method 'name'. Did you forget to initialize an attribute in __init__()?:
File "test-class2.py", line 10
def fn(a: MyClass):
return a.name
~~~~~~ <--- HERE
列舉型別¶
與自定義類類似，列舉型別的語義是使用者定義的，並且整個類定義必須可編譯為 TorchScript IR 並遵守 TorchScript 型別檢查規則。
TSEnumDef ::= "class" Identifier "(enum.Enum | TSEnumType)" ":"
( MemberIdentifier "=" Value )+
( MethodDefinition )*
其中
值必須是型別為
int,float或str的 TorchScript 字面量，並且必須是相同的 TorchScript 型別。TSEnumType是 TorchScript 列舉型別的名稱。與 Python enum 類似，TorchScript 允許受限的Enum子類化，即只有當一個列舉類沒有定義任何成員時，才允許對其進行子類化。
與 Python 相比
TorchScript 只支援
enum.Enum。它不支援其他變體，例如enum.IntEnum,enum.Flag,enum.IntFlag和enum.auto。TorchScript 列舉成員的值必須是相同型別，並且只能是
int,float或str型別，而 Python 列舉成員可以是任何型別。在 TorchScript 中，包含方法的列舉會被忽略。
示例 1
以下示例將類 Color 定義為 Enum 型別
import torch
from enum import Enum
class Color(Enum):
RED = 1
GREEN = 2
def enum_fn(x: Color, y: Color) -> bool:
if x == Color.RED:
return True
return x == y
m = torch.jit.script(enum_fn)
print("Eager: ", enum_fn(Color.RED, Color.GREEN))
print("TorchScript: ", m(Color.RED, Color.GREEN))
示例 2
以下示例展示了受限列舉子類化的情形，其中 BaseColor 未定義任何成員，因此可以被 Color 子類化
import torch
from enum import Enum
class BaseColor(Enum):
def foo(self):
pass
class Color(BaseColor):
RED = 1
GREEN = 2
def enum_fn(x: Color, y: Color) -> bool:
if x == Color.RED:
return True
return x == y
m = torch.jit.script(enum_fn)
print("TorchScript: ", m(Color.RED, Color.GREEN))
print("Eager: ", enum_fn(Color.RED, Color.GREEN))
TorchScript 模組類¶
TSModuleType 是一種特殊的類型別，它從 TorchScript 外部建立的物件例項推斷而來。TSModuleType 的名稱即為物件例項的 Python 類名。Python 類的 __init__() 方法不被視為 TorchScript 方法，因此無需遵守 TorchScript 的型別檢查規則。
模組例項類的型別模式直接從例項物件（在 TorchScript 範圍之外建立）構建，而不是像自定義類那樣從 __init__() 推斷。相同例項類型別的兩個物件可能遵循兩種不同的型別模式。
從這個意義上說，TSModuleType 並不是真正的靜態型別。因此，出於型別安全考慮，TSModuleType 不能用於 TorchScript 型別標註，也不能與 TSType 組合使用。
模組例項類¶
TorchScript 模組型別表示使用者定義的 PyTorch 模組例項的型別模式。當 script 一個 PyTorch 模組時，模組物件總是在 TorchScript 外部建立（即作為引數傳遞給 forward）。Python 模組類被視為模組例項類，因此 Python 模組類的 __init__() 方法不受 TorchScript 型別檢查規則的約束。
TSModuleType ::= "class" Identifier "(torch.nn.Module)" ":"
ClassBodyDefinition
其中
forward()以及用@torch.jit.export裝飾的其他方法必須可編譯為 TorchScript IR 並遵守 TorchScript 的型別檢查規則。
與自定義類不同，只需模組型別的 forward 方法以及用 @torch.jit.export 或 forward() 裝飾的其他方法可編譯即可。最值得注意的是，__init__() 不被視為 TorchScript 方法。因此，模組型別建構函式不能在 TorchScript 範圍內呼叫。相反，TorchScript 模組物件總是在外部構建並傳遞給 torch.jit.script(ModuleObj)。
示例 1
此示例說明了模組型別的一些特性
TestModule例項是在 TorchScript 範圍之外建立的（即在呼叫torch.jit.script之前）。__init__()不被視為 TorchScript 方法，因此無需標註，並且可以包含任意 Python 程式碼。此外，例項類的__init__()方法不能在 TorchScript 程式碼中呼叫。由於TestModule例項是在 Python 中例項化的，在此示例中，TestModule(2.0)和TestModule(2)建立了兩個例項，其資料屬性具有不同的型別。TestModule(2.0)的self.x型別是float，而TestModule(2)的self.y型別是int。TorchScript 會自動編譯透過
@torch.jit.export標註的方法或forward()方法呼叫的其他方法（例如mul()）。TorchScript 程式的入口點可以是模組型別的
forward()、標註為torch.jit.script的函式，或標註為torch.jit.export的方法。
import torch
class TestModule(torch.nn.Module):
def __init__(self, v):
super().__init__()
self.x = v
def forward(self, inc: int):
return self.x + inc
m = torch.jit.script(TestModule(1))
print(f"First instance: {m(3)}")
m = torch.jit.script(TestModule(torch.ones([5])))
print(f"Second instance: {m(3)}")
上面的示例產生以下輸出
First instance: 4
Second instance: tensor([4., 4., 4., 4., 4.])
示例 2
以下示例展示了模組型別的不正確用法。具體來說，此示例在 TorchScript 範圍內呼叫了 TestModule 的建構函式
import torch
class TestModule(torch.nn.Module):
def __init__(self, v):
super().__init__()
self.x = v
def forward(self, x: int):
return self.x + x
class MyModel:
def __init__(self, v: int):
self.val = v
@torch.jit.export
def doSomething(self, val: int) -> int:
# error: should not invoke the constructor of module type
myModel = TestModule(self.val)
return myModel(val)
# m = torch.jit.script(MyModel(2)) # Results in below RuntimeError
# RuntimeError: Could not get name of python class object
型別標註¶
由於 TorchScript 是靜態型別的，程式設計師需要在 TorchScript 程式碼的*關鍵點*標註型別，以便每個區域性變數或例項資料屬性都具有靜態型別，並且每個函式和方法都具有靜態型別簽名。
何時標註型別¶
通常，型別標註只在無法自動推斷靜態型別的地方需要（例如引數或有時是方法或函式的返回型別）。區域性變數和資料屬性的型別通常可以從它們的賦值語句中自動推斷。有時推斷的型別可能過於嚴格，例如，透過賦值 x = None 將 x 推斷為 NoneType，而實際上 x 被用作 Optional。在這種情況下，可能需要型別標註來覆蓋自動推斷，例如 x: Optional[int] = None。請注意，即使可以自動推斷型別，為區域性變數或資料屬性標註型別也是安全的。標註的型別必須與 TorchScript 的型別檢查一致。
當引數、區域性變數或資料屬性未標註型別且無法自動推斷其型別時，TorchScript 假定其預設型別為 TensorType, List[TensorType] 或 Dict[str, TensorType]。
標註函式簽名¶
由於引數可能無法從函式體（包括函式和方法）自動推斷，因此需要進行型別標註。否則，它們將假定預設型別為 TensorType。
TorchScript 支援兩種方法和函式簽名型別註解的風格
Python3 風格直接在簽名上註解型別。因此，它允許某些引數不進行註解（其型別將是
TensorType的預設型別），或者允許不註解返回型別（其型別將自動推斷）。
Python3Annotation ::= "def" Identifier [ "(" ParamAnnot* ")" ] [ReturnAnnot] ":"
FuncOrMethodBody
ParamAnnot ::= Identifier [ ":" TSType ] ","
ReturnAnnot ::= "->" TSType
請注意，在使用 Python3 風格時，型別 self 會自動推斷，不應進行註解。
Mypy 風格將型別註解為緊跟在函式/方法宣告下方的一條註釋。在 Mypy 風格中，由於引數名不出現在註解中，因此所有引數都必須進行註解。
MyPyAnnotation ::= "# type:" "(" ParamAnnot* ")" [ ReturnAnnot ]
ParamAnnot ::= TSType ","
ReturnAnnot ::= "->" TSType
示例 1
在此示例中
a未註解，並假定其預設型別為TensorType。b註解為int型別。返回型別未註解，並自動推斷為
TensorType型別（基於返回值的型別）。
import torch
def f(a, b: int):
return a+b
m = torch.jit.script(f)
print("TorchScript:", m(torch.ones([6]), 100))
示例 2
以下示例使用 Mypy 風格註解。請注意，即使引數或返回值假定為預設型別，也必須進行註解。
import torch
def f(a, b):
# type: (torch.Tensor, int) → torch.Tensor
return a+b
m = torch.jit.script(f)
print("TorchScript:", m(torch.ones([6]), 100))
註解變數和資料屬性¶
通常，資料屬性（包括類資料屬性和例項資料屬性）和區域性變數的型別可以從賦值語句中自動推斷。但是，有時如果變數或屬性與不同型別的值相關聯（例如，None 或 TensorType），則可能需要將其顯式註解為更“寬泛”的型別，例如 Optional[int] 或 Any。
區域性變數¶
區域性變數可以根據 Python3 typing 模組的註解規則進行註解，即，
LocalVarAnnotation ::= Identifier [":" TSType] "=" Expr
通常，區域性變數的型別可以自動推斷。然而在某些情況下，對於可能與不同具體型別相關聯的區域性變數，您可能需要註解一個多型別。典型的多型別包括 Optional[T] 和 Any。
示例
import torch
def f(a, setVal: bool):
value: Optional[torch.Tensor] = None
if setVal:
value = a
return value
ones = torch.ones([6])
m = torch.jit.script(f)
print("TorchScript:", m(ones, True), m(ones, False))
例項資料屬性¶
對於 ModuleType 類，例項資料屬性可以根據 Python3 typing 模組的註解規則進行註解。例項資料屬性可以（可選地）透過 Final 註解為 final。
"class" ClassIdentifier "(torch.nn.Module):"
InstanceAttrIdentifier ":" ["Final("] TSType [")"]
...
其中
InstanceAttrIdentifier是例項屬性的名稱。Final表示該屬性不能在__init__之外重新賦值或在子類中被覆蓋。
示例
import torch
class MyModule(torch.nn.Module):
offset_: int
def __init__(self, offset):
self.offset_ = offset
...
型別註解 API¶
torch.jit.annotate(T, expr)¶
該 API 將型別 T 註解到表示式 expr 上。當表示式的預設型別不是程式設計師預期的型別時，經常使用此方法。例如，空列表（字典）的預設型別是 List[TensorType]（Dict[TensorType, TensorType]），但有時可能用於初始化其他型別的列表。另一個常見用例是註解 tensor.tolist() 的返回型別。然而請注意，它不能用於在 __init__ 中註解模組屬性的型別；在這種情況下應使用 torch.jit.Attribute。
示例
在此示例中，透過 torch.jit.annotate 將 [] 宣告為整數列表（而不是假定 [] 的預設型別為 List[TensorType]）
import torch
from typing import List
def g(l: List[int], val: int):
l.append(val)
return l
def f(val: int):
l = g(torch.jit.annotate(List[int], []), val)
return l
m = torch.jit.script(f)
print("Eager:", f(3))
print("TorchScript:", m(3))
有關更多資訊，請參閱 torch.jit.annotate()。
型別註解附錄¶
TorchScript 型別系統定義¶
TSAllType ::= TSType | TSModuleType
TSType ::= TSMetaType | TSPrimitiveType | TSStructuralType | TSNominalType
TSMetaType ::= "Any"
TSPrimitiveType ::= "int" | "float" | "double" | "complex" | "bool" | "str" | "None"
TSStructuralType ::= TSTuple | TSNamedTuple | TSList | TSDict | TSOptional |
TSUnion | TSFuture | TSRRef | TSAwait
TSTuple ::= "Tuple" "[" (TSType ",")* TSType "]"
TSNamedTuple ::= "namedtuple" "(" (TSType ",")* TSType ")"
TSList ::= "List" "[" TSType "]"
TSOptional ::= "Optional" "[" TSType "]"
TSUnion ::= "Union" "[" (TSType ",")* TSType "]"
TSFuture ::= "Future" "[" TSType "]"
TSRRef ::= "RRef" "[" TSType "]"
TSAwait ::= "Await" "[" TSType "]"
TSDict ::= "Dict" "[" KeyType "," TSType "]"
KeyType ::= "str" | "int" | "float" | "bool" | TensorType | "Any"
TSNominalType ::= TSBuiltinClasses | TSCustomClass | TSEnum
TSBuiltinClass ::= TSTensor | "torch.device" | "torch.stream"|
"torch.dtype" | "torch.nn.ModuleList" |
"torch.nn.ModuleDict" | ...
TSTensor ::= "torch.tensor" and subclasses
不支援的 Typing 構造¶
TorchScript 不支援 Python3 typing 模組的所有特性和型別。本文件中未明確說明的 typing 模組功能均不受支援。下表總結了 TorchScript 中不受支援或受限支援的 typing 構造。
專案 |
描述 |
|
開發中 |
|
不支援 |
|
不支援 |
|
不支援 |
|
不支援 |
|
支援模組屬性、類屬性和註解，但不支援函式。 |
|
不支援 |
|
開發中 |
類型別名 |
不支援 |
名義型別 |
開發中 |
結構型別 |
不支援 |
NewType |
不支援 |
泛型 |
不支援 |
表示式¶
以下部分描述了 TorchScript 中支援的表示式語法。它是參照 Python 語言參考手冊的表示式章節建模的。
算術轉換¶
TorchScript 中執行多項隱式型別轉換
如果一個
Tensor的資料型別為float或int，並且其大小為 0，require_grad未設定為True，且不需要窄化，則可以將其隱式轉換為FloatType或IntType的例項。StringType的例項可以隱式轉換為DeviceType。上述兩個要點的隱式轉換規則可以應用於
TupleType的例項，以生成具有適當包含型別的ListType例項。
顯式轉換可以使用接受原始資料型別作為引數的內建函式 float、int、bool 和 str 呼叫，如果使用者自定義型別實現了 __bool__、__str__ 等方法，也可以接受這些使用者自定義型別。
原子¶
原子是表示式中最基本的元素。
atom ::= identifier | literal | enclosure
enclosure ::= parenth_form | list_display | dict_display
識別符號¶
規定 TorchScript 中合法識別符號的規則與其 Python 對應規則相同。
字面量¶
literal ::= stringliteral | integer | floatnumber
字面量的求值會產生一個具有特定值（必要時對浮點數應用近似值）的相應型別的物件。字面量是不可變的，對相同字面量進行多次求值可能會獲得相同的物件或具有相同值的不同物件。stringliteral、integer 和 floatnumber 的定義方式與其 Python 對應部分相同。
帶括號的形式¶
parenth_form ::= '(' [expression_list] ')'
帶括號的表示式列表的求值結果與表示式列表的求值結果相同。如果列表中至少包含一個逗號，則其結果為一個 Tuple；否則，結果為表示式列表中的單個表示式。一對空括號的結果為一個空 Tuple 物件（Tuple[]）。
列表和字典顯示¶
list_comprehension ::= expression comp_for
comp_for ::= 'for' target_list 'in' or_expr
list_display ::= '[' [expression_list | list_comprehension] ']'
dict_display ::= '{' [key_datum_list | dict_comprehension] '}'
key_datum_list ::= key_datum (',' key_datum)*
key_datum ::= expression ':' expression
dict_comprehension ::= key_datum comp_for
列表和字典可以透過顯式列出容器內容或透過提供一組迴圈指令（即，推導式）來計算內容來構造。推導式在語義上等同於使用 for 迴圈並追加到一個正在構建的列表。推導式隱式建立自己的作用域，以確保目標列表的項不會洩漏到外部作用域。在顯式列出容器項的情況下，表示式列表中的表示式從左到右求值。如果在包含 key_datum_list 的 dict_display 中某個鍵被重複，則結果字典使用列表中使用該重複鍵的最右邊資料項的值。
主項¶
primary ::= atom | attributeref | subscription | slicing | call
下標¶
subscription ::= primary '[' expression_list ']'
The primary 必須求值為一個支援下標的物件。
如果 primary 是一個
List、Tuple或str，則表示式列表必須求值為一個整數或切片。如果 primary 是一個
Dict，則表示式列表必須求值為與Dict的鍵型別相同的物件。如果 primary 是一個
ModuleList，則表示式列表必須是一個integer字面量。如果 primary 是一個
ModuleDict，則表示式必須是一個stringliteral。
切片¶
切片選擇 str、Tuple、List 或 Tensor 中的一系列項。切片可以用作賦值或 del 語句中的表示式或目標。
slicing ::= primary '[' slice_list ']'
slice_list ::= slice_item (',' slice_item)* [',']
slice_item ::= expression | proper_slice
proper_slice ::= [expression] ':' [expression] [':' [expression] ]
切片列表中包含多個切片項的切片只能用於求值為 Tensor 型別物件的 primary。
呼叫¶
call ::= primary '(' argument_list ')'
argument_list ::= args [',' kwargs] | kwargs
args ::= [arg (',' arg)*]
kwargs ::= [kwarg (',' kwarg)*]
kwarg ::= arg '=' expression
arg ::= identifier
The primary 必須脫糖或求值為一個可呼叫物件。所有引數表示式都在嘗試呼叫之前進行求值。
冪運算子¶
power ::= primary ['**' u_expr]
冪運算子的語義與內建的 pow 函式（不支援）相同；它計算左運算元的右運算元次冪。它比左側的一元運算子繫結更緊密，但比右側的一元運算子繫結寬鬆；即 -2 ** -3 == -(2 ** (-3))。左、右運算元可以是 int、float 或 Tensor。在標量-張量/張量-標量指數運算的情況下，標量會被廣播，而張量-張量指數運算是按元素進行的，沒有任何廣播。
一元和算術位運算¶
u_expr ::= power | '-' power | '~' power
一元運算子 - 產生其引數的負數。一元運算子 ~ 產生其引數的按位取反。 - 可用於 int、float 以及 int 和 float 型別的 Tensor。 ~ 只能用於 int 和 int 型別的 Tensor。
二元算術運算¶
m_expr ::= u_expr | m_expr '*' u_expr | m_expr '@' m_expr | m_expr '//' u_expr | m_expr '/' u_expr | m_expr '%' u_expr
a_expr ::= m_expr | a_expr '+' m_expr | a_expr '-' m_expr
二元算術運算子可對 Tensor、int 和 float 進行運算。對於張量-張量運算，兩個引數必須具有相同的形狀。對於標量-張量或張量-標量運算，標量通常會被廣播到張量的大小。除法運算只能接受標量作為右側引數，且不支援廣播。@ 運算子用於矩陣乘法，並且只對 Tensor 引數進行運算。乘法運算子（*）可用於列表和整數，以獲得將原始列表重複一定次數的結果。
移位運算¶
shift_expr ::= a_expr | shift_expr ( '<<' | '>>' ) a_expr
這些運算子接受兩個 int 引數、兩個 Tensor 引數，或者一個 Tensor 引數和一個 int 或 float 引數。在所有情況下，右移 n 位定義為向下取整除以 pow(2, n)，左移 n 位定義為乘以 pow(2, n)。當兩個引數都是 Tensor 時，它們必須具有相同的形狀。當一個是標量而另一個是 Tensor 時，標量在邏輯上會被廣播以匹配 Tensor 的大小。
二元位運算¶
and_expr ::= shift_expr | and_expr '&' shift_expr
xor_expr ::= and_expr | xor_expr '^' and_expr
or_expr ::= xor_expr | or_expr '|' xor_expr
運算子 & 計算其引數的按位與，^ 計算按位異或，| 計算按位或。兩個運算元必須是 int 或 Tensor，或者左運算元是 Tensor 且右運算元是 int。當兩個運算元都是 Tensor 時，它們必須具有相同的形狀。當右運算元是 int 且左運算元是 Tensor 時，右運算元在邏輯上會被廣播以匹配 Tensor 的形狀。
比較¶
comparison ::= or_expr (comp_operator or_expr)*
comp_operator ::= '<' | '>' | '==' | '>=' | '<=' | '!=' | 'is' ['not'] | ['not'] 'in'
比較運算產生一個布林值（True 或 False），或者如果其中一個運算元是 Tensor，則產生一個布林 Tensor。比較可以任意鏈式連線，只要它們產生的布林 Tensors 不包含多個元素。a op1 b op2 c ... 等價於 a op1 b and b op2 c and ...。
值比較¶
運算子 <、>、==、>=、<= 和 != 比較兩個物件的值。這兩個物件通常需要是相同的型別，除非它們之間存在隱式型別轉換。如果使用者自定義型別定義了富比較方法（例如 __lt__），則可以進行比較。內建型別的比較方式與 Python 相同
數字進行數學比較。
字串進行字典序比較。
list、tuple和dict只能與同類型的其他list、tuple和dict進行比較，並使用對應元素的比較運算子進行比較。
成員資格測試運算¶
運算子 in 和 not in 測試成員資格。x in s 的求值結果為 True（如果 x 是 s 的成員）或 False（否則）。x not in s 等價於 not x in s。此運算子支援 list、dict 和 tuple，如果使用者自定義型別實現了 __contains__ 方法，也可以使用此運算子。
身份比較¶
對於除 int、double、bool 和 torch.device 之外的所有型別，運算子 is 和 is not 測試物件的身份；當且僅當 x 和 y 是同一個物件時，x is y 為 True。對於所有其他型別，is 等效於使用 == 進行比較。x is not y 產生 x is y 的反結果。
布林運算¶
or_test ::= and_test | or_test 'or' and_test
and_test ::= not_test | and_test 'and' not_test
not_test ::= 'bool' '(' or_expr ')' | comparison | 'not' not_test
使用者自定義物件可以透過實現 __bool__ 方法來自定義其轉換為 bool 的方式。如果運算元為 false，則運算子 not 產生 True，否則產生 False。表示式 x and y 首先求值 x；如果其結果為 False，則返回其值（False）；否則，求值 y 並返回其值（False 或 True）。表示式 x or y 首先求值 x；如果其結果為 True，則返回其值（True）；否則，求值 y 並返回其值（False 或 True）。
條件表示式¶
conditional_expression ::= or_expr ['if' or_test 'else' conditional_expression]
expression ::= conditional_expression
表示式 x if c else y 首先求值條件 c 而不是 x。如果 c 為 True，則求值 x 並返回其值；否則，求值 y 並返回其值。與 if 語句一樣，x 和 y 必須求值為相同型別的值。
表示式列表¶
expression_list ::= expression (',' expression)* [',']
starred_item ::= '*' primary
帶星號的項只能出現在賦值語句的左側，例如 a, *b, c = ...。
簡單語句¶
以下部分描述了 TorchScript 中支援的簡單語句的語法。它是參照 Python 語言參考手冊的簡單語句章節建模的。
表示式語句¶
expression_stmt ::= starred_expression
starred_expression ::= expression | (starred_item ",")* [starred_item]
starred_item ::= assignment_expression | "*" or_expr
賦值語句¶
assignment_stmt ::= (target_list "=")+ (starred_expression)
target_list ::= target ("," target)* [","]
target ::= identifier
| "(" [target_list] ")"
| "[" [target_list] "]"
| attributeref
| subscription
| slicing
| "*" target
增強賦值語句¶
augmented_assignment_stmt ::= augtarget augop (expression_list)
augtarget ::= identifier | attributeref | subscription
augop ::= "+=" | "-=" | "*=" | "/=" | "//=" | "%=" |
"**="| ">>=" | "<<=" | "&=" | "^=" | "|="
註解賦值語句¶
annotated_assignment_stmt ::= augtarget ":" expression
["=" (starred_expression)]
raise 語句¶
raise_stmt ::= "raise" [expression ["from" expression]]
TorchScript 中的 Raise 語句不支援 try\except\finally。
assert 語句¶
assert_stmt ::= "assert" expression ["," expression]
TorchScript 中的 Assert 語句不支援 try\except\finally。
return 語句¶
return_stmt ::= "return" [expression_list]
TorchScript 中的 Return 語句不支援 try\except\finally。
del 語句¶
del_stmt ::= "del" target_list
pass 語句¶
pass_stmt ::= "pass"
print 語句¶
print_stmt ::= "print" "(" expression [, expression] [.format{expression_list}] ")"
break 語句¶
break_stmt ::= "break"
continue 語句：¶
continue_stmt ::= "continue"
複合語句¶
以下部分描述了 TorchScript 中支援的複合語句的語法。本節還強調了 Torchscript 與常規 Python 語句的不同之處。它是參照 Python 語言參考手冊的複合語句章節建模的。
if 語句¶
Torchscript 支援基本的 if/else 和三元 if/else。
基本的 if/else 語句¶
if_stmt ::= "if" assignment_expression ":" suite
("elif" assignment_expression ":" suite)
["else" ":" suite]
elif 語句可以重複任意多次，但必須在 else 語句之前。
三元 if/else 語句¶
if_stmt ::= return [expression_list] "if" assignment_expression "else" [expression_list]
示例 1
一維 tensor 會被提升為 bool
import torch
@torch.jit.script
def fn(x: torch.Tensor):
if x: # The tensor gets promoted to bool
return True
return False
print(fn(torch.rand(1)))
上面的示例產生以下輸出
True
示例 2
多維 tensor 不會被提升為 bool
import torch
# Multi dimensional Tensors error out.
@torch.jit.script
def fn():
if torch.rand(2):
print("Tensor is available")
if torch.rand(4,5,6):
print("Tensor is available")
print(fn())
執行上述程式碼會導致以下 RuntimeError。
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
@torch.jit.script
def fn():
if torch.rand(2):
~~~~~~~~~~~~ <--- HERE
print("Tensor is available")
RuntimeError: Boolean value of Tensor with more than one value is ambiguous
如果條件變數被註解為 final，則根據條件變數的求值結果，求值 true 分支或 false 分支。
示例 3
在此示例中，由於 a 被註解為 final 並設定為 True，因此只求值 True 分支。
import torch
a : torch.jit.final[Bool] = True
if a:
return torch.empty(2,3)
else:
return []
The while 語句¶
while_stmt ::= "while" assignment_expression ":" suite
Torchscript 中不支援 while…else 語句。它會導致 RuntimeError。
The for-in 語句¶
for_stmt ::= "for" target_list "in" expression_list ":" suite
["else" ":" suite]
Torchscript 中不支援 for...else 語句。它會導致 RuntimeError。
示例 1
對元組的 For 迴圈：這些迴圈會展開（unroll），為元組的每個成員生成一個迴圈體。迴圈體必須對每個成員進行正確的型別檢查。
import torch
from typing import Tuple
@torch.jit.script
def fn():
tup = (3, torch.ones(4))
for x in tup:
print(x)
fn()
上面的示例產生以下輸出
3
1
1
1
1
[ CPUFloatType{4} ]
示例 2
對列表的 For 迴圈：對 nn.ModuleList 的 For 迴圈會在編譯時展開迴圈體，對模組列表的每個成員進行處理。
class SubModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.randn(2))
def forward(self, input):
return self.weight + input
class MyModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.mods = torch.nn.ModuleList([SubModule() for i in range(10)])
def forward(self, v):
for module in self.mods:
v = module(v)
return v
model = torch.jit.script(MyModule())
The with 語句¶
The with 語句用於將一個程式碼塊的執行包裝在由上下文管理器定義的方法中。
with_stmt ::= "with" with_item ("," with_item) ":" suite
with_item ::= expression ["as" target]
如果
with語句中包含目標，上下文管理器的__enter__()方法的返回值將賦值給它。與 Python 不同的是，如果因異常退出程式碼塊（suite），其型別、值和追溯資訊不會作為引數傳遞給__exit__()。會提供三個None引數。try、except和finally語句在with塊內不受支援。在
with塊內引發的異常無法被抑制。
tuple 語句¶
tuple_stmt ::= tuple([iterables])
TorchScript 中的可迭代型別包括
Tensors、lists、tuples、dictionaries、strings、torch.nn.ModuleList和torch.nn.ModuleDict。你不能使用這個內建函式將 List 轉換為 Tuple。
將所有輸出解包到 tuple 中
abc = func() # Function that returns a tuple
a,b = func()
getattr 語句¶
getattr_stmt ::= getattr(object, name[, default])
屬性名必須是字串字面值。
不支援模組型別物件（例如，torch._C）。
不支援自定義類物件（例如，torch.classes.*）。
hasattr 語句¶
hasattr_stmt ::= hasattr(object, name)
屬性名必須是字串字面值。
不支援模組型別物件（例如，torch._C）。
不支援自定義類物件（例如，torch.classes.*）。
zip 語句¶
zip_stmt ::= zip(iterable1, iterable2)
引數必須是可迭代物件。
支援兩個外部容器型別相同但長度不同的可迭代物件。
示例 1
兩個可迭代物件必須是相同的容器型別
a = [1, 2] # List
b = [2, 3, 4] # List
zip(a, b) # works
示例 2
此示例失敗是因為可迭代物件是不同的容器型別
a = (1, 2) # Tuple
b = [2, 3, 4] # List
zip(a, b) # Runtime error
執行上述程式碼會導致以下 RuntimeError。
RuntimeError: Can not iterate over a module list or
tuple with a value that does not have a statically determinable length.
示例 3
支援兩個容器型別相同但資料型別不同的可迭代物件
a = [1.3, 2.4]
b = [2, 3, 4]
zip(a, b) # Works
TorchScript 中的可迭代型別包括 Tensors、lists、tuples、dictionaries、strings、torch.nn.ModuleList 和 torch.nn.ModuleDict。
enumerate 語句¶
enumerate_stmt ::= enumerate([iterable])
引數必須是可迭代物件。
TorchScript 中的可迭代型別包括
Tensors、lists、tuples、dictionaries、strings、torch.nn.ModuleList和torch.nn.ModuleDict。
Python 值¶
解析規則¶
給定一個 Python 值時，TorchScript 嘗試透過以下五種不同的方式解析它
- 可編譯的 Python 實現
當 Python 值由一個可以被 TorchScript 編譯的 Python 實現支援時，TorchScript 會編譯並使用底層的 Python 實現。
示例：
torch.jit.Attribute
- Op Python 封裝器
當 Python 值是原生 PyTorch 操作 (op) 的封裝器時，TorchScript 會生成對應的運算子。
示例：
torch.jit._logging.add_stat_value
- Python 物件身份匹配
對於 TorchScript 支援的有限集合的
torch.*API 呼叫（以 Python 值的形式），TorchScript 嘗試將 Python 值與集合中的每個項進行匹配。匹配成功後，TorchScript 會生成對應的
SugaredValue例項，該例項包含這些值的降低（lowering）邏輯。示例：
torch.jit.isinstance()
- 名稱匹配
對於 Python 內建函式和常量，TorchScript 透過名稱識別它們，並建立一個實現其功能的相應
SugaredValue例項。示例：
all()
- 值快照
對於來自未知模組的 Python 值，TorchScript 嘗試捕獲該值的快照，並將其轉換為正在編譯的函式或方法圖中的常量。
示例：
math.pi
支援的 Python 內建函式¶
內建函式 |
支援級別 |
註釋 |
|---|---|---|
|
部分 |
僅支援 |
|
完全 |
|
|
完全 |
|
|
無 |
|
|
部分 |
僅支援 |
|
部分 |
僅支援 |
|
無 |
|
|
無 |
|
|
無 |
|
|
無 |
|
|
部分 |
僅支援 ASCII 字元集。 |
|
完全 |
|
|
無 |
|
|
無 |
|
|
無 |
|
|
完全 |
|
|
無 |
|
|
完全 |
|
|
完全 |
|
|
無 |
|
|
無 |
|
|
無 |
|
|
部分 |
不遵循 |
|
部分 |
不支援手動指定索引。| 不支援格式型別修飾符。 |
|
無 |
|
|
部分 |
屬性名必須是字串字面值。 |
|
無 |
|
|
部分 |
屬性名必須是字串字面值。 |
|
完全 |
|
|
部分 |
僅支援 |
|
完全 |
僅支援 |
|
無 |
|
|
部分 |
不支援 |
|
完全 |
在檢查 |
|
無 |
|
|
無 |
|
|
完全 |
|
|
完全 |
|
|
部分 |
僅支援 ASCII 字元集。 |
|
完全 |
|
|
部分 |
不支援 |
|
無 |
|
|
完全 |
|
|
無 |
|
|
無 |
|
|
部分 |
不支援 |
|
無 |
|
|
無 |
|
|
完全 |
|
|
部分 |
不支援 |
|
完全 |
|
|
部分 |
不支援 |
|
完全 |
|
|
部分 |
它只能用於 |
|
無 |
|
|
無 |
|
|
完全 |
|
|
無 |
torch.* API¶
遠端過程呼叫¶
TorchScript 支援 RPC API 的一個子集，該子集支援在指定的遠端工作節點而非本地執行函式。
具體來說，以下 API 得到完全支援
torch.distributed.rpc.rpc_sync()rpc_sync()發起一個阻塞的 RPC 呼叫，以便在遠端工作節點上執行一個函式。RPC 訊息與 Python 程式碼的執行並行地傳送和接收。其用法和示例的更多詳細資訊可在
rpc_sync()中找到。
torch.distributed.rpc.rpc_async()rpc_async()發起一個非阻塞的 RPC 呼叫，以便在遠端工作節點上執行一個函式。RPC 訊息與 Python 程式碼的執行並行地傳送和接收。其用法和示例的更多詳細資訊可在
rpc_async()中找到。
torch.distributed.rpc.remote()remote.()在工作節點上執行遠端呼叫，並獲得一個遠端引用RRef作為返回值。其用法和示例的更多詳細資訊可在
remote()中找到。
非同步執行¶
TorchScript 使你能夠建立非同步計算任務，以更好地利用計算資源。這是透過支援僅在 TorchScript 內部使用的一系列 API 來實現的
型別標註¶
TorchScript 是靜態型別的。它提供並支援一組實用工具，以幫助標註變數和屬性。
torch.jit.annotate()在 Python 3 風格的型別提示不太奏效的情況下，為 TorchScript 提供型別提示。
一個常見的例子是為像
[]這樣的表示式標註型別。預設情況下，[]被視為List[torch.Tensor]。當需要不同的型別時，你可以使用此程式碼來提示 TorchScript：torch.jit.annotate(List[int], [])。更多詳細資訊可在
annotate()中找到
torch.jit.Attribute常見的用例包括為
torch.nn.Module屬性提供型別提示。由於它們的__init__方法不被 TorchScript 解析，在模組的__init__方法中應使用torch.jit.Attribute而非torch.jit.annotate。更多詳細資訊可在
Attribute()中找到
torch.jit.FinalPython 的
typing.Final的別名。torch.jit.Final僅為向後相容而保留。
超程式設計¶
TorchScript 提供了一組實用工具來促進超程式設計。
torch.jit.is_scripting()返回一個布林值，指示當前程式是否由
torch.jit.script編譯。當在
assert或if語句中使用時，torch.jit.is_scripting()求值為False的作用域或分支不會被編譯。它的值可以在編譯時靜態求值，因此常用於
if語句中，以阻止 TorchScript 編譯其中一個分支。更多詳細資訊和示例可在
is_scripting()中找到
torch.jit.is_tracing()返回一個布林值，指示當前程式是否由
torch.jit.trace/torch.jit.trace_module跟蹤。更多詳細資訊可在
is_tracing()中找到
@torch.jit.ignore此裝飾器指示編譯器，一個函式或方法應被忽略，並保留為 Python 函式。
這使得你可以在模型中保留尚未與 TorchScript 相容的程式碼。
如果從 TorchScript 呼叫一個由
@torch.jit.ignore裝飾的函式，被忽略的函式會將呼叫分派給 Python 直譯器。包含被忽略函式的模型無法匯出。
更多詳細資訊和示例可在
ignore()中找到
@torch.jit.unused此裝飾器指示編譯器，一個函式或方法應被忽略，並替換為引發異常。
這使得你可以在模型中保留尚未與 TorchScript 相容的程式碼，並且仍然可以匯出模型。
如果從 TorchScript 呼叫一個由
@torch.jit.unused裝飾的函式，將會引發執行時錯誤。更多詳細資訊和示例可在
unused()中找到
型別細化¶
torch.jit.isinstance()返回一個布林值，指示變數是否為指定型別。
其用法和示例的更多詳細資訊可在
isinstance()中找到。
