捷徑

TorchScript 語言參考

本參考手冊描述了 TorchScript 語言的語法和核心語意。TorchScript 是 Python 語言的靜態類型子集。本文檔解釋了 TorchScript 中支援的 Python 功能,以及該語言如何與常規 Python 不同。本參考手冊中未提及的任何 Python 功能都不是 TorchScript 的一部分。TorchScript 專注於在 PyTorch 中表示神經網路模型所需的 Python 功能。

術語

本文档使用以下術語

模式

註解

::=

表示给定的符号被定义为。

" "

代表屬於語法的真實關鍵字和分隔符。

A | B

表示 A 或 B。

( )

表示分組。

[]

表示可選。

A+

表示一個正規表達式,其中術語 A 至少重複一次。

A*

表示一個正規表達式,其中術語 A 重複零次或多次。

類型系統

TorchScript 是 Python 的靜態類型子集。 TorchScript 和完整的 Python 語言之間最大的區別在於 TorchScript 僅支持表達神經網絡模型所需的一小部分類型。

TorchScript 類型

TorchScript 類型系統包含 TSTypeTSModuleType,如下所定義。

TSAllType ::= TSType | TSModuleType
TSType    ::= TSMetaType | TSPrimitiveType | TSStructuralType | TSNominalType

TSType 表示大多數 TorchScript 類型,這些類型是可組合的,並且可以用於 TorchScript 類型註釋。 TSType 指的是以下任何一項

  • Meta 類型,例如 Any

  • 基本類型,例如 intfloatstr

  • 結構類型,例如 Optional[int]List[MyClass]

  • 名義類型 (Python 類),例如 MyClass (用户定义)、torch.tensor (內建)

TSModuleType 表示 torch.nn.Module 及其子類。 它與 TSType 的處理方式不同,因為它的類型模式部分從物件實例推斷,部分從類別定義推斷。 因此,TSModuleType 的實例可能不遵循相同的靜態類型模式。 由於類型安全性的考慮,TSModuleType 不能用作 TorchScript 類型註釋,也不能與 TSType 組合。

Meta 類型

Meta 類型非常抽象,它們更像是類型約束而不是具體類型。 目前,TorchScript 定義了一種 meta 類型,Any,它表示任何 TorchScript 類型。

Any 類型

Any 類型表示任何 TorchScript 類型。 Any 不指定任何類型約束,因此 Any 上沒有類型檢查。 因此,它可以绑定到任何 Python 或 TorchScript 數據類型(例如,int、TorchScript tuple 或未編寫腳本的任意 Python 類別)。

TSMetaType ::= "Any"

其中

  • Any 是 typing 模組中的 Python 類別名稱。 因此,要使用 Any 類型,您必須從 typing 導入它(例如,from typing import Any)。

  • 由於 Any 可以表示任何 TorchScript 類型,因此允許對 Any 上的此類型值進行操作的運算符集受到限制。

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 支持的運算符和方法的一個子集。

設計註解

當我們編寫 PyTorch 模組的腳本時,我們可能會遇到未參與腳本執行的数据。 然而,它必須由類型模式描述。 為未使用的數據(在腳本的上下文中)描述靜態類型不僅麻煩,而且還可能導致不必要的腳本編寫失敗。 引入 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)) 在第二次调用中将元组绑定到 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"

其中

  • TupleListOptionalUnionFutureDict 代表 Python 型別類別名稱,這些名稱定義在 typing 模組中。 若要使用這些型別名稱,您必須從 typing 匯入它們 (例如,from typing import Tuple)。

  • namedtuple 代表 Python 類別 collections.namedtupletyping.NamedTuple

  • FutureRRef 代表 Python 類別 torch.futurestorch.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)))

執行以上程式碼會產生以下指令碼錯誤

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

其中,TSCustomClassTSEnum 必須可編譯為 TorchScript 中繼表示 (IR)。 這是由型別檢查器強制執行的。

內建類別

內建名義型別是 Python 類別,其語意已建置到 TorchScript 系統中 (例如,張量型別)。 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.ModuleListtorch.nn.ModuleDict 在 Python 中定義為清單和字典,但它們在 TorchScript 中的行為更像元組。

  • 在 TorchScript 中,torch.nn.ModuleListtorch.nn.ModuleDict 的實例是不可變的。

  • 迭代 torch.nn.ModuleListtorch.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 中,新型類別是透過繼承物件來指定的。

  • 實例資料屬性是靜態型別的,並且實例屬性必須透過在 __init__() 方法內進行賦值來宣告。

  • 不支援方法多載 (也就是說,您不能有多個具有相同方法名稱的方法)。

  • MethodDefinition 必須可編譯為 TorchScript IR,並遵守 TorchScript 的型別檢查規則 (也就是說,所有方法都必須是有效的 TorchScript 函式,並且類別屬性定義必須是有效的 TorchScript 陳述式)。

  • torch.jit.ignoretorch.jit.unused 可用於忽略無法完全 TorchScript 化或應由編譯器忽略的方法或函式。

與 Python 相比

與 Python 對應物相比,TorchScript 自訂類別受到相當大的限制。 Torchscript 自訂類別

  • 不支援類別屬性。

  • 除了繼承介面型別或物件之外,不支援子類別化。

  • 不支援方法多載。

  • 必須在 __init__() 中初始化其所有實例屬性;這是因為 TorchScript 透過推斷 __init__() 中的屬性型別來建構類別的靜態結構描述。

  • 必須僅包含滿足 TorchScript 型別檢查規則並可編譯為 TorchScript IR 的方法。

範例 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 )*

其中

  • 值必須是 intfloatstr 類型的 TorchScript 常值,且必須是相同的 TorchScript 類型。

  • TSEnumType 是 TorchScript 列舉類型的名稱。與 Python enum 類似,TorchScript 允許受限制的 Enum 子類別化,也就是說,只有在沒有定義任何成員的情況下,才允許列舉子類別化。

與 Python 相比

  • TorchScript 僅支援 enum.Enum。它不支援其他變體,例如 enum.IntEnumenum.Flagenum.IntFlagenum.auto

  • TorchScript 列舉成員的值必須是相同類型,且只能是 intfloatstr 類型,而 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 模組實例的類型架構。 在編寫 PyTorch 模組的腳本時,模組物件始終在 TorchScript 外部建立(即,作為參數傳遞給 forward)。 Python 模組類別被視為模組實例類別,因此 Python 模組類別的 __init__() 方法不受 TorchScript 的類型檢查規則的約束。

TSModuleType ::= "class" Identifier "(torch.nn.Module)" ":"
                    ClassBodyDefinition

其中

  • forward() 和其他使用 @torch.jit.export 裝飾的方法必須可編譯為 TorchScript IR,並受 TorchScript 的類型檢查規則的約束。

與自訂類別不同,只有正向方法和其他使用 @torch.jit.export 裝飾的模組類型的方法需要可編譯。 最值得注意的是,__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.0)self.y 的類型為 int

  • TorchScript 會自動編譯由透過 @torch.jit.exportforward() 方法註釋的方法調用的其他方法(例如,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 會假設其為預設類型 TensorTypeList[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))

註解變數和資料屬性

一般來說,資料屬性(包括類別和實例資料屬性)和區域變數的類型可以從賦值語句中自動推斷。但是,有時,如果變數或屬性與不同類型的值相關聯(例如,NoneTensorType),則可能需要將它們明確地類型註解為更廣泛的類型,例如 Optional[int]Any

區域變數

可以根據 Python3 類型模組註解規則註解區域變數,即:

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 類型模組註解規則註解實例資料屬性。可以選擇性地使用 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。當運算式的預設類型不是程式設計師期望的類型時,通常會使用此 API。例如,一個空列表(字典)的預設類型為 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

不支援的類型建構 (Unsupported Typing Constructs)

TorchScript 並不支援 Python3 typing 模組的所有功能和類型。 本文件中未明確指定的 typing 模組的任何功能均不受支援。 下表總結了 TorchScript 中不支援或支援受限的 typing 建構。

項目 (Item)

描述 (Description)

typing.Any

開發中 (In development)

typing.NoReturn

不支援 (Not supported)

typing.Callable

不支援 (Not supported)

typing.Literal

不支援 (Not supported)

typing.ClassVar

不支援 (Not supported)

typing.Final

支援用於模組屬性、類別屬性和註解,但不支援用於函式。(Supported for module attributes, class attribute, and annotations, but not for functions.)

typing.AnyStr

不支援 (Not supported)

typing.overload

開發中 (In development)

類型別名 (Type aliases)

不支援 (Not supported)

名義類型 (Nominal typing)

開發中 (In development)

結構類型 (Structural typing)

不支援 (Not supported)

NewType

不支援 (Not supported)

泛型 (Generics)

不支援 (Not supported)

運算式 (Expressions)

以下章節描述了 TorchScript 中支援的運算式語法。它以 Python 語言參考的運算式章節為範本。

算術轉換 (Arithmetic Conversions)

在 TorchScript 中執行了許多隱式類型轉換 (There are a number of implicit type conversions that are performed in TorchScript)

  • 具有 floatint 資料類型的 Tensor 可以隱式轉換為 FloatTypeIntType 的實例,前提是它的大小為 0,未將 require_grad 設定為 True,並且不需要縮小範圍。(A Tensor with a float or int data type can be implicitly converted to an instance of FloatType or IntType provided that it has a size of 0, does not have require_grad set to True, and will not require narrowing.)

  • StringType 的實例可以隱式轉換為 DeviceType。(Instances of StringType can be implicitly converted to DeviceType.)

  • 以上兩個要點中的隱式轉換規則可以應用於 TupleType 的實例,以產生具有適當包含類型的 ListType 的實例。(The implicit conversion rules from the two bullet points above can be applied to instances of TupleType to produce instances of ListType with the appropriate contained type.)

可以使用 floatintboolstr 內建函式來調用顯式轉換,這些函式接受原始資料類型作為引數,並且如果它們實現了 __bool____str__ 等,則可以接受使用者定義的類型。(Explicit conversions can be invoked using the float, int, bool, and str built-in functions that accept primitive data types as arguments and can accept user-defined types if they implement __bool__, __str__, etc.)

原子 (Atoms)

原子是運算式的最基本元素。(Atoms are the most basic elements of expressions.)

atom      ::=  identifier | literal | enclosure
enclosure ::=  parenth_form | list_display | dict_display

識別符 (Identifiers)

在 TorchScript 中,決定什麼是合法識別符的規則與其 Python 對應物相同。(The rules that dictate what is a legal identifier in TorchScript are the same as their Python counterparts.)

字面值 (Literals)

literal ::=  stringliteral | integer | floatnumber

字面值的求值會產生具有特定類型和值的物件(並根據需要對浮點數應用近似值)。字面值是不可變的,並且對相同字面值的多次求值可能會獲得相同的物件或具有相同值的不同物件。stringliteralintegerfloatnumber 的定義方式與其 Python 對應物相同。(Evaluation of a literal yields an object of the appropriate type with the specific value (with approximations applied as necessary for floats). Literals are immutable, and multiple evaluations of identical literals may obtain the same object or distinct objects with the same value. stringliteral, integer, and floatnumber are defined in the same way as their Python counterparts.)

括號形式 (Parenthesized Forms)

parenth_form ::=  '(' [expression_list] ')'

括在括號中的運算式清單產生運算式清單產生的任何內容。如果清單包含至少一個逗號,則產生一個 Tuple;否則,產生運算式清單中的單個運算式。一對空的括號產生一個空的 Tuple 物件 (Tuple[])。(A parenthesized expression list yields whatever the expression list yields. If the list contains at least one comma, it yields a Tuple; otherwise, it yields the single expression inside the expression list. An empty pair of parentheses yields an empty Tuple object (Tuple[]).)

清單和字典顯示 (List and Dictionary Displays)

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_listdict_display 中重複一個鍵,則產生的字典使用清單中使用的重複鍵的最右邊資料的值。(Lists and dicts can be constructed by either listing the container contents explicitly or by providing instructions on how to compute them via a set of looping instructions (i.e. a comprehension). A comprehension is semantically equivalent to using a for loop and appending to an ongoing list. Comprehensions implicitly create their own scope to make sure that the items of the target list do not leak into the enclosing scope. In the case that container items are explicitly listed, the expressions in the expression list are evaluated left-to-right. If a key is repeated in a dict_display that has a key_datum_list, the resultant dictionary uses the value from the rightmost datum in the list that uses the repeated key.)

主體 (Primaries)

primary ::=  atom | attributeref | subscription | slicing | call

屬性參考 (Attribute References)

attributeref ::=  primary '.' identifier

primary 必須求值為支援具有名為 identifier 的屬性的屬性參考的類型的物件。(The primary must evaluate to an object of a type that supports attribute references that have an attribute named identifier.)

下標 (Subscriptions)

subscription ::=  primary '[' expression_list ']'

primary 必須求值為支援下標的物件。(The primary must evaluate to an object that supports subscription.)

  • 如果主體是 ListTuplestr,則運算式清單必須求值為整數或切片。(If the primary is a List, Tuple, or str, the expression list must evaluate to an integer or slice.)

  • 如果主體是 Dict,則運算式清單必須求值為與 Dict 的鍵類型相同的物件。(If the primary is a Dict, the expression list must evaluate to an object of the same type as the key type of the Dict.)

  • 如果主體是 ModuleList,則運算式清單必須是 integer 字面值。(If the primary is a ModuleList, the expression list must be an integer literal.)

  • 如果主體是 ModuleDict,則運算式必須是 stringliteral。(If the primary is a ModuleDict, the expression must be a stringliteral.)

切片 (Slicings)

切片 (Slicing) 用於選取 strTupleListTensor 中的一個項目範圍。 切片可以用作表達式,也可以用作賦值或 del 語句中的目標。

slicing      ::=  primary '[' slice_list ']'
slice_list   ::=  slice_item (',' slice_item)* [',']
slice_item   ::=  expression | proper_slice
proper_slice ::=  [expression] ':' [expression] [':' [expression] ]

切片列表中有多個切片項目的切片,只能用於計算結果為 Tensor 類型物件的主體。

呼叫 (Calls)

call          ::=  primary '(' argument_list ')'
argument_list ::=  args [',' kwargs] | kwargs
args          ::=  [arg (',' arg)*]
kwargs        ::=  [kwarg (',' kwarg)*]
kwarg         ::=  arg '=' expression
arg           ::=  identifier

primary 必須去糖 (desugar) 或計算結果為可呼叫物件。 所有引數表達式都會在嘗試呼叫之前進行計算。

冪運算子 (Power Operator)

power ::=  primary ['**' u_expr]

冪運算子的語義與內建的 pow 函數相同 (不支援);它計算其左引數的右引數次方。 它的綁定優先級高於左側的一元運算子,但低於右側的一元運算子;也就是說, -2 ** -3 == -(2 ** (-3))。 左側和右側運算元可以是 intfloatTensor。 在純量-張量/張量-純量指數運算的情況下,純量會被廣播 (broadcast),而張量-張量指數運算則是逐元素進行,不進行任何廣播。

一元和算術位元運算 (Unary and Arithmetic Bitwise Operations)

u_expr ::=  power | '-' power | '~' power

一元 - 運算子產生其引數的否定。 一元 ~ 運算子產生其引數的按位反轉。 - 可用於 intfloatintfloatTensor~ 只能與 intintTensor 一起使用。

二元算術運算 (Binary Arithmetic Operations)

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

二元算術運算子可以對 Tensorintfloat 進行運算。 對於張量-張量運算,兩個引數必須具有相同的形狀 (shape)。 對於純量-張量或張量-純量運算,純量通常會廣播到張量的大小。 除法運算子只能接受純量作為其右側引數,並且不支援廣播。@ 運算子用於矩陣乘法,並且僅對 Tensor 引數進行運算。 乘法運算子 (*) 可以與列表和整數一起使用,以獲得重複特定次數的原始列表的結果。

位移運算 (Shifting Operations)

shift_expr ::=  a_expr | shift_expr ( '<<' | '>>' ) a_expr

這些運算子接受兩個 int 引數、兩個 Tensor 引數,或者一個 Tensor 引數和一個 intfloat 引數。 在所有情況下,右移 n 定義為底除法 (floor division) 除以 pow(2, n),左移 n 定義為乘以 pow(2, n)。 當兩個引數都是 Tensors 時,它們必須具有相同的形狀。 當一個是純量,另一個是 Tensor 時,純量會邏輯上廣播以匹配 Tensor 的大小。

二元位元運算 (Binary Bitwise Operations)

and_expr ::=  shift_expr | and_expr '&' shift_expr
xor_expr ::=  and_expr | xor_expr '^' and_expr
or_expr  ::=  xor_expr | or_expr '|' xor_expr

& 運算子計算其引數的按位 AND,^ 計算按位 XOR,| 計算按位 OR。 兩個運算元都必須是 intTensor,或者左側運算元必須是 Tensor,右側運算元必須是 int。 當兩個運算元都是 Tensor 時,它們必須具有相同的形狀。 當右側運算元是 int 且左側運算元是 Tensor 時,右側運算元會在邏輯上廣播以匹配 Tensor 的形狀。

比較 (Comparisons)

comparison    ::=  or_expr (comp_operator or_expr)*
comp_operator ::=  '<' | '>' | '==' | '>=' | '<=' | '!=' | 'is' ['not'] | ['not'] 'in'

比較產生一個布林值 (Boolean value) ( TrueFalse ),或者,如果其中一個運算元是 Tensor,則會產生一個布林 Tensor。 只要比較不會產生具有多個元素的布林 Tensors,就可以任意鏈式比較。 a op1 b op2 c ... 等同於 a op1 b and b op2 c and ...

數值比較 (Value Comparisons)

運算子 <>==>=<=!= 比較兩個物件的值。這兩個物件通常需要是相同的類型,除非物件之間存在隱含的類型轉換。如果使用者定義的類型定義了豐富的比較方法(例如,__lt__),則可以進行比較。內建類型比較的工作方式與 Python 類似。

  • 數字會以數學方式進行比較。

  • 字串會以字典順序進行比較。

  • liststuplesdicts 只能與相同類型的其他 liststuplesdicts 進行比較,並使用對應元素的比較運算子進行比較。

成員資格測試運算

運算子 innot in 用於測試成員資格。如果 xs 的成員,則 x in s 的結果為 True,否則為 Falsex not in s 等同於 not x in s。此運算子支援 listsdictstuples,並且如果使用者定義的類型實作了 __contains__ 方法,也可以使用。

身分比較

對於除了 intdoublebooltorch.device 之外的所有類型,運算子 isis not 用於測試物件的身分;當且僅當 xy 是同一個物件時,x is yTrue。對於所有其他類型,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 並傳回其值(FalseTrue)。 運算式 x or y 會先評估 x;如果它是 True,則傳回其值(True);否則,評估 y 並傳回其值(FalseTrue)。

條件運算式

conditional_expression ::=  or_expr ['if' or_test 'else' conditional_expression]
 expression            ::=  conditional_expression

運算式 x if c else y 會先評估條件 c,而不是 x。如果 cTrue,則評估 x 並傳回其值;否則,評估 y 並傳回其值。與 if 語句一樣,xy 必須評估為相同類型的值。

運算式清單

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

具有 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

在此範例中,僅評估 True 分支,因為 a 被註解為 final 並設定為 True

import torch

a : torch.jit.final[Bool] = True

if a:
    return torch.empty(2,3)
else:
    return []

while 語句

while_stmt ::=  "while" assignment_expression ":" suite

Torchscript 不支援 while…else 語句。 這會導致 RuntimeError

for-in 語句

for_stmt ::=  "for" target_list "in" expression_list ":" suite
              ["else" ":" suite]

Torchscript 不支援 for...else 語句。 這會導致 RuntimeError

範例 1

在元組上的 For 迴圈:這些會展開迴圈,為元組的每個成員產生一個主體。 對於每個成員,主體必須正確進行類型檢查。

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

with 語句

with 語句用於使用上下文管理器定義的方法來包裝區塊的執行。

with_stmt ::=  "with" with_item ("," with_item) ":" suite
with_item ::=  expression ["as" target]
  • 如果在 with 語句中包含目標,則上下文管理器的 __enter__() 的傳回值會分配給它。 與 python 不同,如果異常導致套件退出,則其類型、值和追蹤不會作為引數傳遞給 __exit__()。 提供三個 None 引數。

  • with 區塊內不支援 tryexceptfinally 語句。

  • 無法抑制在 with 區塊中引發的例外。

tuple 語句

tuple_stmt ::= tuple([iterables])
  • TorchScript 中的可迭代類型包括 Tensorsliststuplesdictionariesstringstorch.nn.ModuleListtorch.nn.ModuleDict

  • 您無法使用此內建函數將列表轉換為元組。

將所有輸出解壓縮到元組中由涵蓋

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 中的可迭代類型包括 Tensorsliststuplesdictionariesstringstorch.nn.ModuleListtorch.nn.ModuleDict

enumerate 語句

enumerate_stmt ::= enumerate([iterable])
  • 引數必須是可迭代的。

  • TorchScript 中的可迭代類型包含 Tensorsliststuplesdictionariesstringstorch.nn.ModuleListtorch.nn.ModuleDict

Python 值

解析規則

當給定一個 Python 值時,TorchScript 會嘗試以下列五種不同的方式解析它:

  • 可編譯的 Python 實作
    • 當一個 Python 值由 TorchScript 可以編譯的 Python 實作支持時,TorchScript 會編譯並使用底層的 Python 實作。

    • 範例:torch.jit.Attribute

  • Op Python 包裝器
    • 當 Python 值是原生 PyTorch 運算的包裝器時,TorchScript 會發出相應的運算符。

    • 範例:torch.jit._logging.add_stat_value

  • Python 物件識別碼匹配
    • 對於 TorchScript 支援的一組有限的 torch.* API 呼叫(以 Python 值的形式),TorchScript 會嘗試將 Python 值與集合中的每個項目進行匹配。

    • 當匹配時,TorchScript 會產生一個相應的 SugaredValue 實例,其中包含這些值的降低邏輯。

    • 範例:torch.jit.isinstance()

  • 名稱匹配
    • 對於 Python 內建函數和常數,TorchScript 會按名稱識別它們,並創建一個相應的 SugaredValue 實例來實現它們的功能。

    • 範例:all()

  • 值快照
    • 對於來自無法識別的模組的 Python 值,TorchScript 會嘗試取得該值的快照,並將其轉換為正在編譯的函數或方法圖中的常數。

    • 範例:math.pi

Python 內建函數支援

TorchScript 對 Python 內建函數的支援

內建函數

支援程度

註解

abs()

部分

僅支援 Tensor/Int/Float 類型輸入。 | 不會採用 __abs__ 覆寫。

all()

完整

any()

完整

ascii()

bin()

部分

僅支援 Int 類型輸入。

bool()

部分

僅支援 Tensor/Int/Float 類型輸入。

breakpoint()

bytearray()

bytes()

callable()

chr()

部分

僅支援 ASCII 字元集。

classmethod()

完整

compile()

complex()

delattr()

dict()

完整

dir()

divmod()

完整

enumerate()

完整

eval()

exec()

filter()

float()

部分

不會採用 __index__ 覆寫。

format()

部分

不支援手動索引指定。 | 不支援格式類型修飾符。

frozenset()

getattr()

部分

屬性名稱必須是字串文字。

globals()

hasattr()

部分

屬性名稱必須是字串文字。

hash()

完整

Tensor 的雜湊基於識別碼,而不是數值。

hex()

部分

僅支援 Int 類型輸入。

id()

完整

僅支援 Int 類型輸入。

input()

int()

部分

不支援 base 參數。 | 不會採用 __index__ 覆寫。

isinstance()

完整

當針對容器類型(例如 Dict[str, int])進行檢查時,torch.jit.isintance 提供更好的支援。

issubclass()

iter()

len()

完整

list()

完整

ord()

部分

僅支援 ASCII 字元集。

pow()

完整

print()

部分

不支援 separateendfile 參數。

property()

range()

完整

repr()

reversed()

round()

部分

不支援 ndigits 參數。

set()

setattr()

slice()

完整

sorted()

部分

不支援 key 參數。

staticmethod()

完整

str()

部分

不支援 encodingerrors 參數。

sum()

完整

super()

部分

它只能在 nn.Module__init__ 方法中使用。

type()

vars()

zip()

完整

__import__()

Python 內建值支援

TorchScript 對 Python 內建值的支援

內建值

支援程度

註解

False

完整

True

完整

完整

NotImplemented

Ellipsis

完整

torch.* APIs

遠端程序呼叫

TorchScript 支援 RPC API 的一個子集,該子集支援在指定的遠端工作者上執行函數,而不是在本機執行。

具體來說,完全支援以下 API:

  • torch.distributed.rpc.rpc_sync()
    • rpc_sync() 進行阻塞式 RPC 呼叫,以在遠端工作者上執行函數。 RPC 訊息的傳送和接收與 Python 程式碼的執行並行。

    • 有關其用法和範例的更多詳細資訊,請參閱 rpc_sync()

  • torch.distributed.rpc.rpc_async()
    • rpc_async() 會對遠端 worker 發出非阻擋式的 RPC 呼叫,以執行函式。RPC 訊息的傳送和接收會與 Python 程式碼的執行並行處理。

    • 關於其用法和範例的更多詳細資訊,請參閱 rpc_async()

  • torch.distributed.rpc.remote()
    • remote.() 在 worker 上執行遠端呼叫,並取得 Remote Reference RRef 作為回傳值。

    • 關於其用法和範例的更多詳細資訊,請參閱 remote()

非同步執行

TorchScript 讓您可以建立非同步計算任務,以更有效地利用計算資源。 這是透過支援一系列僅在 TorchScript 中可用的 API 來完成的。

  • torch.jit.fork()
    • 建立一個非同步任務,執行 func,並回傳此執行結果值的參考。 Fork 將立即返回。

    • 等同於 torch.jit._fork(),僅為了向後相容性而保留。

    • 關於其用法和範例的更多詳細資訊,請參閱 fork()

  • torch.jit.wait()
    • 強制完成 torch.jit.Future[T] 非同步任務,並回傳任務的結果。

    • 等同於 torch.jit._wait(),僅為了向後相容性而保留。

    • 關於其用法和範例的更多詳細資訊,請參閱 wait()

類型註解

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.Final
    • Python 的 typing.Final 的別名。torch.jit.Final 僅為了向後相容性而保留。

元編程

TorchScript 提供了一組實用工具來促進元編程

  • torch.jit.is_scripting()
    • 回傳一個布林值,指示目前的程式是否由 torch.jit.script 編譯。

    • 當在 assertif 語句中使用時,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()

文件

訪問 PyTorch 的完整開發者文檔

查看文檔

教程

獲取針對初學者和高級開發者的深入教程

查看教程

資源

查找開發資源並獲得您問題的解答

查看資源