捷徑

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 指的是以下任何一種:

  • 中繼類型,例如 Any

  • 基本類型,例如 intfloatstr

  • 結構類型,例如 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 或未編寫腳本的任意 Python 類別)。

TSMetaType ::= "Any"

其中

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

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

Any 類型支援的運算子

  • 賦值給 Any 類型的資料。

  • 綁定到 Any 類型的參數或返回值。

  • x isx is not,其中 x 的類型為 Any

  • isinstance(x, Type),其中 x 的類型為 Any

  • 類型為 Any 的資料可供列印。

  • 如果 List[Any] 類型的資料是相同類型 T 的值清單,並且 T 支援比較運算子,則該資料可能是可排序的。

與 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) 將浮點數類型繫結到 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 表示在模組 typing 中定義的 Python 類型類別名稱。若要使用這些類型名稱,您必須從 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)。這是由類型檢查器強制執行的。

內建類別

內建標稱類型是其語義內建於 TorchScript 系統中的 Python 類別(例如,張量類型)。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 列舉類似,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 的類型檢查規則的約束。

與自訂類別不同,只有模組類型的 forward 方法和其他使用 @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.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 = Nonex 推斷為 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 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(可選)註釋為最終的。

"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"

TSStructualType ::=  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

不支援的類型建構

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

項目

說明

typing.Any

開發中

typing.NoReturn

不支援

typing.Callable

不支援

typing.Literal

不支援

typing.ClassVar

不支援

typing.Final

支援模組屬性、類別屬性和標註,但不支援函式。

typing.AnyStr

不支援

typing.overload

開發中

類型別名

不支援

名義類型

開發中

結構類型

不支援

NewType

不支援

泛型

不支援

表達式

以下章節說明 TorchScript 中支援的表達式語法。它仿照 Python 語言參考的表達式章節

算術轉換

TorchScript 中會執行許多隱式類型轉換

  • 資料類型為 floatintTensor 可以隱式轉換為 FloatTypeIntType 的實例,前提是它的大小為 0,require_grad 未設定為 True,並且不需要縮小。

  • StringType 的實例可以隱式轉換為 DeviceType

  • 以上兩點中的隱式轉換規則可以應用於 TupleType 的實例,以產生具有適當包含類型的 ListType 的實例。

可以使用 floatintboolstr 內建函數來叫用顯式轉換,這些函數接受原始資料類型作為參數,並且如果它們實現了 __bool____str__ 等,則可以接受使用者定義的類型。

原子

原子是表達式的最基本元素。

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

識別字

規定 TorchScript 中合法識別字的規則与其 Python 對應項 相同。

字面值

literal ::=  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_listdict_display 中重複金鑰,則產生的字典會使用列表中最右側使用重複金鑰的資料中的值。

主體

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

屬性參考

attributeref ::=  primary '.' identifier

primary 必須評估為支援具有名為 identifier 的屬性的屬性參考的類型的物件。

註標

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

primary 必須評估為支援註標的物件。

  • 如果主體是 ListTuplestr,則表達式列表必須評估為整數或切片。

  • 如果主體是 Dict,則表達式列表必須評估為與 Dict 的金鑰類型相同的類型的物件。

  • 如果主體是 ModuleList,則表達式列表必須是 integer 字面值。

  • 如果主體是 ModuleDict,則表達式必須是 stringliteral

切片

切片會在 strTupleListTensor 中選擇一系列項目。切片可以在賦值或 del 陳述式中用作表達式或目標。

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

切片列表中包含多個切片項目的切片只能與評估為類型 Tensor 的物件的主體一起使用。

呼叫

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

primary 必須反糖化或評估為可呼叫物件。所有參數表達式都會在嘗試呼叫之前進行評估。

冪運算符

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

冪運算符的語義與內建 pow 函數(不支援)相同;它會計算其左運算元提高到其右運算元的冪。它比左側的一元運算符綁定得更緊密,但比右側的一元運算符綁定得更鬆散;即 -2 ** -3 == -(2 ** (-3))。左運算元和右運算元可以是 intfloatTensor。在純量-張量/張量-純量冪運算的情況下,純量會廣播,而張量-張量冪運算是逐元素完成的,没有任何廣播。

一元和算術位元運算

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

一元 - 運算符會產生其參數的否定。一元 ~ 運算符會產生其參數的位元反轉。 - 可以與 intfloatintfloat 類型的 Tensor 一起使用。 ~ 只能與 intint 類型的 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

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

移位運算

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

這些運算子接受兩個 int 引數、兩個 Tensor 引數,或一個 Tensor 引數和一個 intfloat 引數。在所有情況下,向右移位 n 位定義為除以 pow(2, n) 的結果向下取整,向左移位 n 位定義為乘以 pow(2, n)。當兩個引數都是 Tensors 時,它們必須具有相同的形狀。當一個是純量而另一個是 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

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

比較

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

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

值比較

運算子 <>==>=<=!= 比較兩個物件的值。這兩個物件通常需要是同一種類型,除非物件之間存在隱式類型轉換。如果在使用者定義的類型上定義了豐富的比較方法(例如 __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 測試物件的身分;x is y 僅在 xy 是同一個物件時才為 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 並返回其值(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

具有一維的 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

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

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

針對 lists 的 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

  • 您無法使用此內建函數將清單轉換為 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 中的可迭代類型包括 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 op 的包裝器時,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.* API

遠端程序呼叫

TorchScript 支援一部分 RPC API,這些 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 來完成的

  • 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 的完整開發人員文件

檢視文件

教學

取得適用於初學者和進階開發人員的深入教學

檢視教學

資源

尋找開發資源並取得問題解答

檢視資源