TorchScript 语言参考¶
本参考手册描述了TorchScript语言的语法和核心语义。 TorchScript是Python语言的一个静态类型子集。本文档解释了TorchScript中支持的Python特性，以及该语言与常规Python的不同之处。任何未在本参考手册中提到的Python特性都不属于TorchScript。TorchScript特别关注Python中用于在PyTorch中表示神经网络模型的特性。
术语¶
本文档使用以下术语：
模式 |
注释 |
|---|---|
|
指示给定的符号被定义为。 |
|
表示作为语法一部分的实际关键字和分隔符。 |
|
表示 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 类型¶
The 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类型。任何类型的数据都是可打印的。
类型为
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"
其中：
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
# 错误：元组未被识别，因为未从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)))
运行上述代码会产生以下脚本错误：
文件 "test-tuple.py", 行 5, 在 <模块>
def inc(x: Tuple[int, int]):
名称错误: 名称 'Tuple' 未定义
解决方法是添加一行 from typing import Tuple 到代码的开头。
名义类型¶
名义上的 TorchScript 类型是 Python 类。这些类型被称为名义类型，因为它们是通过自定义名称声明的，并且使用类名进行比较。名义类进一步分为以下几类：
TSNominalType ::= TSBuiltinClasses | TSCustomClass | TSEnum
其中，TSCustomClass 和 TSEnum 必须可编译为 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.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" ]
"类" ClassName [ "(对象)" ] ":"
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相比
TorchScript 自定义类与其 Python 对应类相比，功能相当有限。TorchScript 自定义类：
不支持类属性。
不支持子类化，除非是子类化接口类型或对象。
不支持方法重载。
必须在
__init__()中初始化所有实例属性；这是因为 TorchScript 通过推断__init__()中的属性类型来构建类的静态模式。必须仅包含符合 TorchScript 类型检查规则且可编译为 TorchScript IR 的方法。
示例 1
Python 类可以在 TorchScript 中使用，如果它们被标注为 @torch.jit.script，类似于声明一个 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
# 错误: self.x 未在 __init__ 中定义
def assign_x(self):
self.x = torch.rand(2, 3)
该类将无法编译并发出以下错误：
运行时错误：
尝试设置不存在的属性：x。你是否忘记在 __init__() 中初始化它？:
def assign_x(self):
self.x = torch.rand(2, 3)
~~~~~~~~~~~~~~~~~~~~~~~~ <--- 此处
示例 3
在这个例子中，一个TorchScript自定义类定义了一个类变量name，这是不允许的：
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' 对象没有属性或方法 'name'。你是否忘记在 __init__() 中初始化一个属性？:
File "test-class2.py", line 10
def fn(a: MyClass):
return a.name
~~~~~~ <--- 此处
枚举类型¶
与自定义类类似，枚举类型的语义是由用户定义的，整个类定义必须可编译为 TorchScript IR 并遵循 TorchScript 类型检查规则。
TSEnumDef ::= "class" Identifier "(enum.Enum | TSEnumType)" ":"
( MemberIdentifier "=" Value )+
( MethodDefinition )*
其中：
值必须是类型为
int、float或str的 TorchScript 字面量，并且必须是相同的 TorchScript 类型。TSEnumType是 TorchScript 枚举类型的名称。类似于 Python 枚举，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 模块实例的类型模式。当脚本化一个 PyTorch 模块时，模块对象总是在 TorchScript 外部创建的（即，作为参数传递给 forward）。Python 模块类被视为模块实例类，因此 Python 模块类的 __init__() 方法不受 TorchScript 类型检查规则的约束。
TSModuleType ::= "类" 标识符 "(torch.nn.Module)" ":"
类体定义
其中：
forward()和其他用@torch.jit.export装饰的方法必须可编译为 TorchScript IR 并遵循 TorchScript 的类型检查规则。
与自定义类不同，只有模块类型的前向方法和其他用 @torch.jit.export 装饰的方法需要可编译。值得注意的是，__init__() 不被视为 TorchScript 方法。因此，模块类型构造函数不能在 TorchScript 的作用域内调用。相反，TorchScript 模块对象总是在外部构造并传递给 torch.jit.script(ModuleObj)。
示例 1
此示例说明了模块类型的一些特性：
在TorchScript的范围之外创建了
TestModule实例（即在调用torch.jit.script之前）。__init__()不被认为是 TorchScript 方法，因此，它不需要被注释并且可以包含任意的 Python 代码。此外，实例类的__init__()方法不能在 TorchScript 代码中被调用。因为TestModule实例是在 Python 中实例化的，在这个例子中，TestModule(2.0)和TestModule(2)创建了两个实例，它们的数据属性具有不同的类型。self.x对于TestModule(2.0)是float类型，而self.y对于TestModule(2.0)是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"第一个实例: {m(3)}")
m = torch.jit.script(TestModule(torch.ones([5])))
print(f"第二个实例: {m(3)}")
上面的示例产生以下输出：
第一个实例: 4
第二个实例: 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:
# 错误: 不应该调用模块类型的构造函数
myModel = TestModule(self.val)
return myModel(val)
# m = torch.jit.script(MyModel(2)) # 结果是下面的RuntimeError
# RuntimeError: 无法获取python类对象的名称
类型注解¶
由于 TorchScript 是静态类型的，程序员需要在 TorchScript 代码的关键点注释类型，以便每个局部变量或实例数据属性都有一个静态类型，并且每个函数和方法都有一个静态类型的签名。
何时注释类型¶
一般来说，类型注解只在静态类型无法自动推断的地方才需要（例如，参数或有时是方法或函数的返回类型）。局部变量和数据属性的类型通常会从它们的赋值语句中自动推断出来。有时推断的类型可能过于严格，例如，x 通过赋值 x = None 被推断为 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):
# 类型: (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注解为最终属性。
"类" ClassIdentifier "(torch.nn.Module):"
InstanceAttrIdentifier ":" ["最终("] 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 构造。
项目 |
描述 |
|
在开发中 |
|
不支持 |
|
不支持 |
|
不支持 |
|
不支持 |
|
支持模块属性、类属性和注解，但不支持函数。 |
|
不支持 |
|
在开发中 |
类型别名 |
不支持 |
名义类型 |
在开发中 |
结构化类型 |
不支持 |
新类型 |
不支持 |
泛型 |
不支持 |
表达式¶
以下部分描述了TorchScript中支持的表达式的语法。 它是基于Python语言参考中的表达式章节建模的。
算术转换¶
在TorchScript中，会执行多种隐式类型转换：
一个具有
float或int数据类型的Tensor可以隐式转换为FloatType或IntType的实例，前提是它的大小为 0，没有将require_grad设置为True，并且不需要缩小。实例
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[]）。
列表和字典显示¶
列表推导式 ::= 表达式 comp_for
comp_for ::= 'for' 目标列表 'in' or_expr
列表显示 ::= '[' [表达式列表 | 列表推导式] ']'
字典显示 ::= '{' [键值对列表 | 字典推导式] '}'
键值对列表 ::= 键值对 (',' 键值对)*
键值对 ::= 表达式 ':' 表达式
字典推导式 ::= 键值对 comp_for
列表和字典可以通过显式列出容器内容或通过提供一组循环指令（即推导式）来构建。推导式在语义上等同于使用for循环并将元素追加到正在进行的列表中。推导式隐式创建自己的作用域，以确保目标列表的项不会泄漏到封闭作用域中。如果容器项是显式列出的，则表达式列表中的表达式从左到右进行求值。如果在具有key_datum_list的dict_display中重复了某个键，则生成的字典将使用使用该重复键的列表中最右侧数据项的值。
主要内容¶
primary ::= atom | attributeref | subscription | slicing | call
属性引用¶
attributeref ::= primary '.' identifier
The primary 必须计算为一个支持属性引用的对象类型，该类型具有名为 identifier 的属性。
订阅¶
subscription ::= primary '[' expression_list ']'
The primary 必须计算为一个支持订阅的对象。
如果主键是
List、Tuple或str，表达式列表必须计算为整数或切片。如果主键是一个
字典，表达式列表必须计算为与字典的键类型相同类型的对象。如果主键是一个
ModuleList，表达式列表必须是一个整数字面量。如果主键是一个
ModuleDict，表达式必须是一个字符串字面量。
切片¶
切片选择 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类型的对象。
调用¶
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 和 Tensor 的 int 和 float 一起使用。~ 只能与 int 和 Tensor 的 int 一起使用。
二进制算术运算¶
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)。当两个参数都是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
& 运算符计算其参数的按位与，^ 计算按位异或，| 计算按位或。
两个操作数必须是 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：
数字是按数学方式进行比较的。
字符串按字典顺序进行比较。
lists,tuples, 和dicts只能与相同类型的lists,tuples, 和dicts进行比较，并且使用相应元素的比较运算符进行比较。
成员测试操作¶
运算符 in 和 not in 用于测试成员资格。x in s 如果 x 是 s 的成员，则计算结果为 True，否则为 False。
x not in s 等价于 not x in s。此运算符支持 lists、dicts 和 tuples，并且如果用户定义的类型实现了 __contains__ 方法，也可以与它们一起使用。
身份比较¶
对于所有类型，除了 int、double、bool 和 torch.device，运算符 is 和 is not 用于测试对象的身份；
x is y 当且仅当 x 和 y 是同一个对象时为 True。对于所有其他类型，is 等价于
使用 == 进行比较。x is not y 产生 x is y 的逆结果。
布尔运算¶
or_test ::= and_test | or_test '或' and_test
and_test ::= not_test | and_test '和' not_test
not_test ::= '布尔' '(' or_expr ')' | comparison | '非' not_test
用户定义的对象可以通过实现一个__bool__方法来自定义其转换为bool的行为。运算符not在其操作数为假时返回True，否则返回False。表达式x和y首先计算x；如果它是False，则返回其值（False）；否则，计算y并返回其值（False或True）。表达式x或y首先计算x；如果它是True，则返回其值（True）；否则，计算y并返回其值（False或True）。
条件表达式¶
conditional_expression ::= or_expr ['如果' or_test '否则' 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)]
The raise 语句¶
raise_stmt ::= "raise" [expression ["from" expression]]
TorchScript 中的 raise 语句不支持 try\except\finally。
The assert 语句¶
assert_stmt ::= "assert" expression ["," expression]
TorchScript 中的断言语句不支持 try\except\finally。
The return 语句¶
return_stmt ::= "return" [expression_list]
TorchScript 中的 return 语句不支持 try\except\finally。
The del 语句¶
del_stmt ::= "del" target_list
The pass 语句¶
pass_stmt ::= "pass"
The print 语句¶
print_stmt ::= "print" "(" expression [, expression] [.format{expression_list}] ")"
The break 语句¶
break_stmt ::= "break"
The continue 语句:¶
continue_stmt ::= "继续"
复合语句¶
以下部分描述了TorchScript中支持的复合语句的语法。 本节还重点介绍了TorchScript与常规Python语句的不同之处。 它是根据Python语言参考中的复合语句章节建模的。
The if 语句¶
Torchscript 支持基本的 if/else 和三元 if/else。
基本的 if/else 语句¶
if_stmt ::= "如果" assignment_expression ":" suite
("否则如果" assignment_expression ":" suite)
["其他" ":" suite]
elif 语句可以重复任意次数，但它需要在 else 语句之前。
三元 if/else 语句¶
if_stmt ::= return [expression_list] "如果" assignment_expression "否则" [expression_list]
示例 1
一个 tensor 具有1个维度会被提升为 bool：
import torch
@torch.jit.script
def fn(x: torch.Tensor):
if x: # 张量被提升为布尔类型
return True
return False
print(fn(torch.rand(1)))
上面的示例产生以下输出：
真
示例 2
具有多维的 tensor 不会被提升为 bool：
import torch
# 多维张量会报错。
@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: 以下操作在TorchScript解释器中失败。
TorchScript的追溯信息（最近一次调用最后）：
@torch.jit.script
def fn():
if torch.rand(2):
~~~~~~~~~~~~ <--- 此处
print("Tensor is available")
RuntimeError: 具有多个值的张量的布尔值是不明确的
如果一个条件变量被注释为final，则根据条件变量的评估结果，执行真分支或假分支。
示例 3
在这个例子中，只有True分支被评估，因为a被标注为final并设置为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
while…else 语句在 Torchscript 中不受支持。它会导致 RuntimeError。
The for-in 语句¶
for_stmt ::= "for" target_list "in" expression_list ":" suite
["else" ":" suite]
for...else 语句在Torchscript中不受支持。这会导致一个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())
The with 语句¶
使用 with 语句来包装由上下文管理器定义的方法的块的执行。
with_stmt ::= "with" with_item ("," with_item) ":" suite
with_item ::= expression ["as" target]
如果在
with语句中包含了一个目标，上下文管理器的__enter__()的返回值将被赋值给它。与 Python 不同，如果异常导致代码块退出，异常的类型、值和回溯不会作为参数传递给__exit__()。将提供三个None参数。try,except, 和finally语句在with块中不被支持。在
with块中引发的异常无法被抑制。
The tuple 语句¶
tuple_stmt ::= tuple([iterables])
TorchScript 中的可迭代类型包括
Tensors、列表、元组、字典、字符串、torch.nn.ModuleList和torch.nn.ModuleDict。您不能使用此内置函数将列表转换为元组。
将所有输出解包到元组中，由以下内容涵盖：
abc = func() # 返回一个元组的函数
a,b = func()
The getattr 语句¶
getattr_stmt ::= getattr(object, name[, default])
属性名称必须是一个字面字符串。
不支持模块类型对象（例如，torch._C）。
不支持自定义类对象（例如，torch.classes.*）。
The hasattr 语句¶
hasattr_stmt ::= hasattr(object, name)
属性名称必须是字面字符串。
不支持模块类型对象（例如，torch._C）。
不支持自定义类对象（例如，torch.classes.*）。
The zip 语句¶
zip_stmt ::= zip(iterable1, iterable2)
参数必须是可迭代对象。
支持两个外部容器类型相同但长度不同的可迭代对象。
示例 1
两个可迭代对象必须是相同的容器类型：
a = [1, 2] # 列表
b = [2, 3, 4] # 列表
zip(a, b) # 有效
示例 2
这个示例失败了，因为可迭代对象属于不同的容器类型：
a = (1, 2) # 元组
b = [2, 3, 4] # 列表
zip(a, b) # 运行时错误
运行上述代码会产生以下RuntimeError。
RuntimeError: 无法 迭代 一个 模块 列表 或
元组 带有 一个 值 不 具有 静态 确定 长度.
示例 3
支持相同容器类型但不同数据类型的两个可迭代对象：
a = [1.3, 2.4]
b = [2, 3, 4]
zip(a, b) # 有效
TorchScript 中的可迭代类型包括 Tensors、列表、元组、字典、字符串、torch.nn.ModuleList 和 torch.nn.ModuleDict。
The enumerate 语句¶
enumerate_stmt ::= enumerate([iterable])
参数必须是可迭代对象。
TorchScript 中的可迭代类型包括
Tensors、列表、元组、字典、字符串、torch.nn.ModuleList和torch.nn.ModuleDict。
Python 值¶
解析规则¶
当给定一个Python值时，TorchScript尝试通过以下五种不同的方式来解析它：
- Compilable Python Implementation:
当一个Python值由一个可以被TorchScript编译的Python实现支持时，TorchScript会编译并使用底层Python实现。
示例:
torch.jit.Attribute
- Op Python Wrapper:
当一个Python值是原生PyTorch操作的包装器时，TorchScript会发出相应的操作符。
示例:
torch.jit._logging.add_stat_value
- Python Object Identity Match:
对于TorchScript支持的一组有限的
torch.*API调用（以Python值的形式），TorchScript尝试将Python值与集合中的每个项进行匹配。当匹配时，TorchScript 生成一个相应的
SugaredValue实例，其中包含这些值的降低逻辑。示例:
torch.jit.isinstance()
- Name Match:
对于Python内置函数和常量，TorchScript通过名称识别它们，并创建一个相应的
SugaredValue实例来实现它们的功能。示例:
all()
- Value Snapshot:
对于来自未识别模块的Python值，TorchScript尝试对这些值进行快照，并将其转换为正在编译的函数或方法的图中的常量。
示例:
math.pi
Python 内置函数支持¶
内置函数 |
支持级别 |
注释 |
|---|---|---|
|
部分 |
仅支持 |
|
完整 |
|
|
完整 |
|
|
无 |
|
|
部分 |
仅支持 |
|
部分 |
仅支持 |
|
无 |
|
|
无 |
|
|
无 |
|
|
无 |
|
|
部分 |
仅支持ASCII字符集。 |
|
完整 |
|
|
无 |
|
|
无 |
|
|
无 |
|
|
完整 |
|
|
无 |
|
|
完整 |
|
|
完整 |
|
|
无 |
|
|
无 |
|
|
无 |
|
|
部分 |
不遵循 |
|
部分 |
不支持手动索引指定。 | 不支持格式类型修饰符。 |
|
无 |
|
|
部分 |
属性名称必须是字符串字面量。 |
|
无 |
|
|
部分 |
属性名称必须是字符串字面量。 |
|
完整 |
|
|
部分 |
仅支持 |
|
完整 |
仅支持 |
|
无 |
|
|
部分 |
|
|
完整 |
|
|
无 |
|
|
无 |
|
|
完整 |
|
|
完整 |
|
|
部分 |
仅支持ASCII字符集。 |
|
完整 |
|
|
部分 |
|
|
无 |
|
|
完整 |
|
|
无 |
|
|
无 |
|
|
部分 |
|
|
无 |
|
|
无 |
|
|
完整 |
|
|
部分 |
|
|
完整 |
|
|
部分 |
|
|
完整 |
|
|
部分 |
它只能在 |
|
无 |
|
|
无 |
|
|
完整 |
|
|
无 |
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.()在 worker 上执行远程调用并获得一个远程引用RRef作为返回值。更多关于其用法和示例的详细信息可以在
remote()中找到。
异步执行¶
TorchScript 使您能够创建异步计算任务，以更好地利用计算资源。这是通过支持一系列仅在 TorchScript 中可用的 API 来实现的：
类型注解¶
TorchScript 是静态类型的。它提供并支持一组工具来帮助注释变量和属性：
torch.jit.annotate()为TorchScript提供类型提示，因为在Python 3风格类型提示效果不佳的情况下。
一个常见的例子是为表达式
[]注释类型。默认情况下，[]被视为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兼容的代码。
如果一个由
@torch.jit.ignore装饰的函数在 TorchScript 中被调用，被忽略的函数将会将调用分派给 Python 解释器。带有忽略函数的模型无法导出。
更多详情和示例可以在
ignore()中找到
@torch.jit.unused这个装饰器向编译器指示应忽略一个函数或方法，并将其替换为引发异常。
这允许你在模型中保留尚未与TorchScript兼容的代码，并且仍然可以导出你的模型。
如果一个由
@torch.jit.unused装饰的函数在 TorchScript 中被调用，将会引发一个运行时错误。更多详情和示例可以在
unused()中找到
类型细化¶
torch.jit.isinstance()返回一个布尔值，指示变量是否为指定类型。
关于其用法和示例的更多详细信息可以在
isinstance()中找到。