TorchScript 语言参考

TorchScript 是 Python 的一个静态类型子集，它既可以直接编写（使用 @torch.jit.script 装饰器），也可以通过跟踪从 Python 代码自动生成。使用跟踪时，代码会通过仅记录张量上的实际算子并简单执行和丢弃其他周围 Python 代码自动转换为此 Python 子集。

当使用 @torch.jit.script 装饰器直接编写 TorchScript 时，程序员只能使用 TorchScript 中支持的 Python 子集。本节记录了 TorchScript 中支持的内容，就像它是独立语言的语言参考一样。本参考中未提及的 Python 任何功能都不属于 TorchScript。有关可用的 PyTorch 张量方法、模块和函数的完整参考，请参阅 内置函数。

作为 Python 的子集，任何有效的 TorchScript 函数也是一个有效的 Python 函数。这使得可以禁用 TorchScript并使用标准 Python 工具（如 pdb）调试函数。反之则不然：有许多有效的 Python 程序不是有效的 TorchScript 程序。相反，TorchScript 专门关注 PyTorch 中表示神经网络模型所需的 Python 特性。

类型

TorchScript 与完整 Python 语言之间最大的区别在于，TorchScript 仅支持表示神经网络模型所需的一小部分类型。具体来说，TorchScript 支持

类型	描述
张量	具有任意数据类型、维度或后端的 PyTorch 张量
Tuple[T0, T1, ..., TN]	包含子类型 T0、T1 等的元组（例如 Tuple[Tensor, Tensor]）
bool	布尔值
int	标量整数
float	标量浮点数
str	字符串
List[T]	其所有成员均为类型 T 的列表
Optional[T]	值为 None 或类型 T 的值
Dict[K, V]	键类型为 K 且值类型为 V 的字典。仅允许 str、int 和 float 作为键类型。
T	TorchScript 类
E	TorchScript 枚举
NamedTuple[T0, T1, ...]	collections.namedtuple 元组类型
Union[T0, T1, ...]	子类型 T0、T1 等之一

与 Python 不同，TorchScript 函数中的每个变量都必须具有一个静态类型。这使得优化 TorchScript 函数变得更容易。

示例（类型不匹配）

import torch

@torch.jit.script
def an_error(x):
    if x:
        r = torch.rand(1)
    else:
        r = 4
    return r

Traceback (most recent call last):
  ...
RuntimeError: ...

Type mismatch: r is set to type Tensor in the true branch and type int in the false branch:
@torch.jit.script
def an_error(x):
    if x:
    ~~~~~
        r = torch.rand(1)
        ~~~~~~~~~~~~~~~~~
    else:
    ~~~~~
        r = 4
        ~~~~~ <--- HERE
    return r
and was used here:
    else:
        r = 4
    return r
           ~ <--- HERE...

不支持的类型构造

TorchScript 不支持 typing 模块的所有特性和类型。其中一些是更基本的东西，不太可能在未来添加，而另一些则可能在有足够的用户需求使其成为优先事项时添加。

以下类型和特性来自 typing 模块，在 TorchScript 中不可用。

项目	描述
typing.Any	typing.Any 目前正在开发中，但尚未发布
typing.NoReturn	未实现
typing.Sequence	未实现
typing.Callable	未实现
typing.Literal	未实现
typing.ClassVar	未实现
typing.Final	这适用于 模块属性 类属性注释，但不适用于函数
typing.AnyStr	TorchScript 不支持 bytes，因此不使用此类型
typing.overload	typing.overload 目前正在开发中，但尚未发布
类型别名	未实现
名义子类型化与结构子类型化	名义类型化正在开发中，但结构类型化没有
NewType	不太可能实现
泛型	不太可能实现

typing 模块中未在此文档中明确列出的任何其他功能均不受支持。

默认类型

默认情况下，假定 TorchScript 函数的所有参数都是张量。要指定 TorchScript 函数的参数是其他类型，可以使用上面列出的类型使用 MyPy 风格的类型注释。

import torch

@torch.jit.script
def foo(x, tup):
    # type: (int, Tuple[Tensor, Tensor]) -> Tensor
    t0, t1 = tup
    return t0 + t1 + x

print(foo(3, (torch.rand(3), torch.rand(3))))

注意

还可以使用 typing 模块中的 Python 3 类型提示来注释类型。

import torch
from typing import Tuple

@torch.jit.script
def foo(x: int, tup: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
    t0, t1 = tup
    return t0 + t1 + x

print(foo(3, (torch.rand(3), torch.rand(3))))

假定空列表是 List[Tensor]，空字典是 Dict[str, Tensor]。要实例化其他类型的空列表或字典，请使用 Python 3 类型提示。

示例（Python 3 的类型注释）

import torch
import torch.nn as nn
from typing import Dict, List, Tuple

class EmptyDataStructures(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> Tuple[List[Tuple[int, float]], Dict[str, int]]:
        # This annotates the list to be a `List[Tuple[int, float]]`
        my_list: List[Tuple[int, float]] = []
        for i in range(10):
            my_list.append((i, x.item()))

        my_dict: Dict[str, int] = {}
        return my_list, my_dict

x = torch.jit.script(EmptyDataStructures())

可选类型细化

当在 if 语句的条件中对变量类型 Optional[T] 与 None 进行比较或在 assert 中进行检查时，TorchScript 将优化该变量的类型。编译器可以推断出多个 None 检查，这些检查与 and、or 和 not 结合使用。未显式编写的 if 语句的 else 块也会进行优化。

必须在 if 语句的条件中进行 None 检查；将 None 检查分配给变量并在 if 语句的条件中使用它不会优化检查中变量的类型。只有局部变量会被优化，像 self.x 这样的属性不会被优化，并且必须分配给局部变量才能被优化。

示例（优化参数和局部变量上的类型）

import torch
import torch.nn as nn
from typing import Optional

class M(nn.Module):
    z: Optional[int]

    def __init__(self, z):
        super().__init__()
        # If `z` is None, its type cannot be inferred, so it must
        # be specified (above)
        self.z = z

    def forward(self, x, y, z):
        # type: (Optional[int], Optional[int], Optional[int]) -> int
        if x is None:
            x = 1
            x = x + 1

        # Refinement for an attribute by assigning it to a local
        z = self.z
        if y is not None and z is not None:
            x = y + z

        # Refinement via an `assert`
        assert z is not None
        x += z
        return x

module = torch.jit.script(M(2))
module = torch.jit.script(M(None))

TorchScript 类

警告

TorchScript 类支持处于实验阶段。目前它最适合简单的记录类型（比如带有附加方法的 NamedTuple）。

如果 Python 类用 @torch.jit.script 进行注释，则可以在 TorchScript 中使用它们，类似于声明 TorchScript 函数的方式

@torch.jit.script
class Foo:
  def __init__(self, x, y):
    self.x = x

  def aug_add_x(self, inc):
    self.x += inc

此子集受到限制

• 所有函数都必须是有效的 TorchScript 函数（包括 __init__()）。

• 类必须是新式类，因为我们使用 __new__() 用 pybind11 构造它们。

• TorchScript 类是静态类型的。成员只能在 __init__() 方法中分配给 self 来声明。

  例如，在 __init__() 方法之外分配给 self

  @torch.jit.script
  class Foo:
    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
  

• 类的主体中不允许出现除方法定义之外的表达式。

• 不支持继承或任何其他多态性策略，除了从 object 继承以指定一个新式类。

定义一个类后，它可以在 TorchScript 和 Python 中交替使用，就像任何其他 TorchScript 类型一样。

# Declare a TorchScript class
@torch.jit.script
class Pair:
  def __init__(self, first, second):
    self.first = first
    self.second = second

@torch.jit.script
def sum_pair(p):
  # type: (Pair) -> Tensor
  return p.first + p.second

p = Pair(torch.rand(2, 3), torch.rand(2, 3))
print(sum_pair(p))

TorchScript 枚举

Python 枚举可以在 TorchScript 中使用，无需任何额外的注释或代码。

from enum import Enum


class Color(Enum):
    RED = 1
    GREEN = 2

@torch.jit.script
def enum_fn(x: Color, y: Color) -> bool:
    if x == Color.RED:
        return True

    return x == y

定义枚举后，它可以在 TorchScript 和 Python 中交替使用，就像任何其他 TorchScript 类型一样。枚举值的类型必须是 int、float 或 str。所有值必须是同一种类型；不支持枚举值的异构类型。

命名元组

由 collections.namedtuple 生成的类型可以在 TorchScript 中使用。

import torch
import collections

Point = collections.namedtuple('Point', ['x', 'y'])

@torch.jit.script
def total(point):
    # type: (Point) -> Tensor
    return point.x + point.y

p = Point(x=torch.rand(3), y=torch.rand(3))
print(total(p))

可迭代对象

某些函数（例如，zip 和 enumerate）只能对可迭代类型进行操作。TorchScript 中的可迭代类型包括 Tensor、列表、元组、字典、字符串、torch.nn.ModuleList 和 torch.nn.ModuleDict。

表达式

支持以下 Python 表达式。

字面量

True
False
None
'string literals'
"string literals"
3  # interpreted as int
3.4  # interpreted as a float

列表构造

空列表被假定具有 List[Tensor] 类型。其他列表文本字面的类型由成员的类型派生。有关更多详细信息，请参阅 默认类型。

[3, 4]
[]
[torch.rand(3), torch.rand(4)]

元组构造

(3, 4)
(3,)

字典构造

空字典被假定具有 Dict[str, Tensor] 类型。其他字典文本字面的类型由成员的类型派生。有关更多详细信息，请参阅 默认类型。

{'hello': 3}
{}
{'a': torch.rand(3), 'b': torch.rand(4)}

变量

有关如何解析变量，请参阅 变量解析。

my_variable_name

算术运算符

a + b
a - b
a * b
a / b
a ^ b
a @ b

比较运算符

a == b
a != b
a < b
a > b
a <= b
a >= b

逻辑运算符

a and b
a or b
not b

下标和切片

t[0]
t[-1]
t[0:2]
t[1:]
t[:1]
t[:]
t[0, 1]
t[0, 1:2]
t[0, :1]
t[-1, 1:, 0]
t[1:, -1, 0]
t[i:j, i]

函数调用

对 内置函数 的调用

torch.rand(3, dtype=torch.int)

对其他脚本函数的调用

import torch

@torch.jit.script
def foo(x):
    return x + 1

@torch.jit.script
def bar(x):
    return foo(x)

方法调用

对内置类型（如张量）的方法的调用：x.mm(y)

在模块上，必须先编译方法，然后才能调用它们。TorchScript 编译器在编译其他方法时会递归编译它看到的方法。默认情况下，编译从 forward 方法开始。由 forward 调用的任何方法都将被编译，由这些方法调用的任何方法也将被编译，依此类推。要从 forward 以外的方法开始编译，请使用 @torch.jit.export 装饰器（forward 隐式标记为 @torch.jit.export）。

直接调用子模块（例如 self.resnet(input)）等效于调用其 forward 方法（例如 self.resnet.forward(input)）。

import torch
import torch.nn as nn
import torchvision

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        means = torch.tensor([103.939, 116.779, 123.68])
        self.means = torch.nn.Parameter(means.resize_(1, 3, 1, 1))
        resnet = torchvision.models.resnet18()
        self.resnet = torch.jit.trace(resnet, torch.rand(1, 3, 224, 224))

    def helper(self, input):
        return self.resnet(input - self.means)

    def forward(self, input):
        return self.helper(input)

    # Since nothing in the model calls `top_level_method`, the compiler
    # must be explicitly told to compile this method
    @torch.jit.export
    def top_level_method(self, input):
        return self.other_helper(input)

    def other_helper(self, input):
        return input + 10

# `my_script_module` will have the compiled methods `forward`, `helper`,
# `top_level_method`, and `other_helper`
my_script_module = torch.jit.script(MyModule())

三元表达式

x if x > y else y

强制转换

float(ten)
int(3.5)
bool(ten)
str(2)``

访问模块参数

self.my_parameter
self.my_submodule.my_parameter

语句

TorchScript 支持以下类型的语句

简单赋值

a = b
a += b # short-hand for a = a + b, does not operate in-place on a
a -= b

模式匹配赋值

a, b = tuple_or_list
a, b, *c = a_tuple

多重赋值

a = b, c = tup

打印语句

print("the result of an add:", a + b)

If 语句

if a < 4:
    r = -a
elif a < 3:
    r = a + a
else:
    r = 3 * a

除了布尔值，浮点数、整数和张量也可以用于条件中，并且将隐式转换为布尔值。

While 循环

a = 0
while a < 4:
    print(a)
    a += 1

带有范围的 For 循环

x = 0
for i in range(10):
    x *= i

遍历元组的 For 循环

这些循环展开循环，为元组的每个成员生成一个主体。主体必须针对每个成员进行正确类型检查。

tup = (3, torch.rand(4))
for x in tup:
    print(x)

遍历常量 nn.ModuleList 的 For 循环

要在编译方法中使用 nn.ModuleList，必须通过将属性的名称添加到类型的 __constants__ 列表中来将其标记为常量。遍历 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):
    __constants__ = ['mods']

    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


m = torch.jit.script(MyModule())

Break 和 Continue

for i in range(5):
    if i == 1:
        continue
    if i == 3:
        break
    print(i)

Return

return a, b

变量解析

TorchScript 支持 Python 的变量解析（即作用域）规则的一个子集。局部变量的行为与 Python 中相同，但有一个限制，即变量在函数的所有路径中必须具有相同的类型。如果变量在 if 语句的不同分支中具有不同的类型，则在 if 语句结束之后使用它会出错。

类似地，如果变量仅在函数的某些路径中定义，则不允许使用该变量。

示例

@torch.jit.script
def foo(x):
    if x < 0:
        y = 4
    print(y)

Traceback (most recent call last):
  ...
RuntimeError: ...

y is not defined in the false branch...
@torch.jit.script...
def foo(x):
    if x < 0:
    ~~~~~~~~~
        y = 4
        ~~~~~ <--- HERE
    print(y)
and was used here:
    if x < 0:
        y = 4
    print(y)
          ~ <--- HERE...

在函数定义时，非局部变量在编译时解析为 Python 值。然后，使用 Python 值的使用 中描述的规则将这些值转换为 TorchScript 值。

Python 值的使用

为了让编写 TorchScript 更加方便，我们允许脚本代码引用周围作用域中的 Python 值。例如，任何时候引用 torch，TorchScript 编译器实际上在函数声明时将其解析为 torch Python 模块。这些 Python 值不是 TorchScript 的一等部分。相反，它们在编译时被反糖化为 TorchScript 支持的基本类型。这取决于编译发生时引用的 Python 值的动态类型。本节描述了在 TorchScript 中访问 Python 值时使用的规则。

函数

TorchScript 可以调用 Python 函数。当逐步将模型转换为 TorchScript 时，此功能非常有用。模型可以逐个函数移动到 TorchScript，同时保留对 Python 函数的调用。这样，你可以在进行的过程中逐步检查模型的正确性。

torch.jit.is_scripting()

在编译时返回 True，否则返回 False 的函数。这在使用 @unused 装饰器时特别有用，可将模型中尚未与 TorchScript 兼容的代码保留在模型中。.. testcode

import torch

@torch.jit.unused
def unsupported_linear_op(x):
    return x

def linear(x):
   if torch.jit.is_scripting():
      return torch.linear(x)
   else:
      return unsupported_linear_op(x)

返回类型

bool

torch.jit.is_tracing()

返回布尔值。

在跟踪（如果在使用 torch.jit.trace 跟踪代码期间调用函数）时返回 True，否则返回 False。

Python 模块上的属性查找

TorchScript 可以查找模块上的属性。torch.add 等内置函数通过这种方式进行访问。这允许 TorchScript 调用在其他模块中定义的函数。

Python 定义的常量

TorchScript 还提供了一种使用在 Python 中定义的常量的方法。这些常量可用于将超参数硬编码到函数中，或定义通用常量。有两种方法可以指定应将 Python 值视为常量。

1. 将模块属性中查找的值视为常量

import math
import torch

@torch.jit.script
def fn():
    return math.pi

2. 可以通过使用 Final[T] 进行注释，将 ScriptModule 的属性标记为常量

import torch
import torch.nn as nn

class Foo(nn.Module):
    # `Final` from the `typing_extensions` module can also be used
    a : torch.jit.Final[int]

    def __init__(self):
        super().__init__()
        self.a = 1 + 4

    def forward(self, input):
        return self.a + input

f = torch.jit.script(Foo())

支持的 Python 常量类型包括

• int

• float

• bool

• torch.device

• torch.layout

• torch.dtype

• 包含受支持类型的元组

• torch.nn.ModuleList，可用于 TorchScript for 循环

模块属性

可以使用 torch.nn.Parameter 包装器和 register_buffer 将张量分配给模块。如果可以推断出其类型，则将分配给已编译模块的其他值添加到已编译模块中。TorchScript 中可用的所有 类型 都可用作模块属性。张量属性在语义上与缓冲区相同。无法推断出空列表和字典以及 None 值的类型，必须通过 PEP 526 样式 类注释指定。如果无法推断出类型并且未明确注释，则不会将其作为属性添加到生成的 ScriptModule 中。

示例

from typing import List, Dict

class Foo(nn.Module):
    # `words` is initialized as an empty list, so its type must be specified
    words: List[str]

    # The type could potentially be inferred if `a_dict` (below) was not
    # empty, but this annotation ensures `some_dict` will be made into the
    # proper type
    some_dict: Dict[str, int]

    def __init__(self, a_dict):
        super().__init__()
        self.words = []
        self.some_dict = a_dict

        # `int`s can be inferred
        self.my_int = 10

    def forward(self, input):
        # type: (str) -> int
        self.words.append(input)
        return self.some_dict[input] + self.my_int

f = torch.jit.script(Foo({'hi': 2}))