TorchScript

• TorchScript 语言参考

TorchScript 是一种从 PyTorch 代码创建可序列化和可优化模型的方法。任何 TorchScript 程序都可以从 Python 进程保存并在没有 Python 依赖项的进程中加载。

我们提供了一些工具，可以将模型从纯 Python 程序逐步转换为可以在独立于 Python 的环境中运行的 TorchScript 程序，例如在独立的 C++ 程序中。这使得可以使用 Python 中熟悉的工具在 PyTorch 中训练模型，然后通过 TorchScript 将模型导出到生产环境，在生产环境中，Python 程序可能由于性能和多线程原因而处于劣势。

有关 TorchScript 的简要介绍，请参阅 TorchScript 简介 教程。

有关将 PyTorch 模型转换为 TorchScript 并在 C++ 中运行的端到端示例，请参阅 在 C++ 中加载 PyTorch 模型 教程。

创建 TorchScript 代码

script	脚本化函数。
trace	跟踪一个函数并返回一个可执行文件或 ScriptFunction，该文件将使用即时编译进行优化。
script_if_tracing	在跟踪期间首次调用 fn 时对其进行编译。
trace_module	跟踪模块并返回一个可执行的 ScriptModule，该模块将使用即时编译进行优化。
fork	创建一个异步任务来执行 func，并创建一个对此执行结果值的引用。
wait	强制完成 torch.jit.Future[T] 异步任务，并返回任务的结果。
ScriptModule	带有方法、属性和参数的 C++ torch::jit::Module 的包装器。
ScriptFunction	功能上等同于 ScriptModule，但表示单个函数，并且没有任何属性或参数。
freeze	冻结 ScriptModule，将子模块和属性内联为常量。
optimize_for_inference	执行一组优化过程，以优化模型以进行推理。
enable_onednn_fusion	根据参数 enabled 启用或禁用 onednn JIT 融合。
onednn_fusion_enabled	返回是否启用了 onednn JIT 融合。
set_fusion_strategy	设置融合期间可能发生的特殊化的类型和数量。
strict_fusion	如果在推理中并非所有节点都已融合，或在训练中进行了符号微分，则给出错误。
save	保存此模块的离线版本，以便在单独的进程中使用。
load	加载之前使用 torch.jit.save 保存的 ScriptModule 或 ScriptFunction。
ignore	此装饰器向编译器指示应忽略某个函数或方法，并将其保留为 Python 函数。
unused	此装饰器向编译器指示应忽略某个函数或方法，并将其替换为引发异常。
interface	装饰以注释不同类型的类或模块。
isinstance	在 TorchScript 中提供容器类型细化。
Attribute	此方法是一个传递函数，它返回 value，主要用于向 TorchScript 编译器指示左侧表达式是具有 type 类型的类实例属性。
annotate	用于在 TorchScript 编译器中指定 the_value 的类型。

混合跟踪和脚本

在许多情况下，跟踪或脚本都是将模型转换为 TorchScript 的更简单方法。可以组合跟踪和脚本以满足模型部分的特定需求。

脚本函数可以调用跟踪函数。当您需要在简单的顺序模型周围使用控制流时，这尤其有用。例如，序列到序列模型的波束搜索通常用脚本编写，但可以调用使用跟踪生成的编码器模块。

示例（在脚本中调用跟踪函数）

import torch

def foo(x, y):
    return 2 * x + y

traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

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

跟踪函数可以调用脚本函数。当模型的一小部分需要一些控制流，即使模型的大部分只是一个前馈网络时，这很有用。跟踪函数调用的脚本函数内部的控制流会被正确保留。

示例（在跟踪函数中调用脚本函数）

import torch

@torch.jit.script
def foo(x, y):
    if x.max() > y.max():
        r = x
    else:
        r = y
    return r


def bar(x, y, z):
    return foo(x, y) + z

traced_bar = torch.jit.trace(bar, (torch.rand(3), torch.rand(3), torch.rand(3)))

这种组合也适用于 nn.Module，其中它可以用于使用跟踪生成一个子模块，该子模块可以从脚本模块的方法中调用。

示例（使用跟踪模块）

import torch
import torchvision

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

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

my_script_module = torch.jit.script(MyScriptModule())

TorchScript 语言

TorchScript 是 Python 的静态类型子集，因此许多 Python 特性直接适用于 TorchScript。有关详细信息，请参阅完整的 TorchScript 语言参考。

内置函数和模块

TorchScript 支持使用大多数 PyTorch 函数和许多 Python 内置函数。有关受支持函数的完整参考，请参阅 TorchScript 内置函数。

PyTorch 函数和模块

TorchScript 支持 PyTorch 提供的张量和神经网络函数的子集。Tensor 上的大多数方法以及 torch 命名空间中的函数、torch.nn.functional 中的所有函数和 torch.nn 中的大多数模块在 TorchScript 中都受支持。

有关不受支持的 PyTorch 函数和模块的列表，请参阅 TorchScript 不支持的 PyTorch 结构。

Python 函数和模块

TorchScript 支持许多 Python 的 内置函数。math 模块也受支持（有关详细信息，请参阅 math 模块），但不支持其他 Python 模块（内置或第三方）。

Python 语言参考比较

有关受支持的 Python 特性的完整列表，请参阅 Python 语言参考覆盖范围。

调试

禁用 JIT 以进行调试

PYTORCH_JIT

设置环境变量 PYTORCH_JIT=0 将禁用所有脚本和跟踪注释。如果您的 TorchScript 模型之一中存在难以调试的错误，则可以使用此标志强制所有内容都使用原生 Python 运行。由于使用此标志禁用了 TorchScript（脚本和跟踪），因此您可以使用 pdb 之类的工具来调试模型代码。例如

@torch.jit.script
def scripted_fn(x : torch.Tensor):
    for i in range(12):
        x = x + x
    return x

def fn(x):
    x = torch.neg(x)
    import pdb; pdb.set_trace()
    return scripted_fn(x)

traced_fn = torch.jit.trace(fn, (torch.rand(4, 5),))
traced_fn(torch.rand(3, 4))

使用 pdb 调试此脚本是可行的，除非我们调用 @torch.jit.script 函数。我们可以全局禁用 JIT，以便我们可以将 @torch.jit.script 函数作为普通的 Python 函数调用，而不是编译它。如果上面的脚本被称为 disable_jit_example.py，我们可以像这样调用它

$ PYTORCH_JIT=0 python disable_jit_example.py

这样我们就可以像普通的 Python 函数一样单步进入 @torch.jit.script 函数。要禁用特定函数的 TorchScript 编译器，请参阅 @torch.jit.ignore。

检查代码

TorchScript 为所有 ScriptModule 实例提供了一个代码美化打印器。此美化打印器将脚本方法的代码解释为有效的 Python 语法。例如

@torch.jit.script
def foo(len):
    # type: (int) -> torch.Tensor
    rv = torch.zeros(3, 4)
    for i in range(len):
        if i < 10:
            rv = rv - 1.0
        else:
            rv = rv + 1.0
    return rv

print(foo.code)

具有单个 forward 方法的 ScriptModule 将具有属性 code，您可以使用它来检查 ScriptModule 的代码。如果 ScriptModule 具有多个方法，则需要访问方法本身的 .code 而不是模块。我们可以通过访问 .foo.code 来检查 ScriptModule 上名为 foo 的方法的代码。上面的示例产生以下输出

def foo(len: int) -> Tensor:
    rv = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
    rv0 = rv
    for i in range(len):
        if torch.lt(i, 10):
            rv1 = torch.sub(rv0, 1., 1)
        else:
            rv1 = torch.add(rv0, 1., 1)
        rv0 = rv1
    return rv0

这是 TorchScript 对 forward 方法代码的编译结果。您可以使用它来确保 TorchScript（跟踪或脚本化）已正确捕获您的模型代码。

解释图

TorchScript 还以 IR 图的形式具有比代码美化打印器更低级别的表示形式。

TorchScript 使用静态单赋值 (SSA) 中间表示 (IR) 来表示计算。这种格式的指令包括 ATen（PyTorch 的 C++ 后端）运算符和其他基本运算符，包括用于循环和条件的控制流运算符。例如

@torch.jit.script
def foo(len):
    # type: (int) -> torch.Tensor
    rv = torch.zeros(3, 4)
    for i in range(len):
        if i < 10:
            rv = rv - 1.0
        else:
            rv = rv + 1.0
    return rv

print(foo.graph)

graph 遵循 检查代码 部分中关于 forward 方法查找的相同规则。

上面的示例脚本生成以下图

graph(%len.1 : int):
  %24 : int = prim::Constant[value=1]()
  %17 : bool = prim::Constant[value=1]() # test.py:10:5
  %12 : bool? = prim::Constant()
  %10 : Device? = prim::Constant()
  %6 : int? = prim::Constant()
  %1 : int = prim::Constant[value=3]() # test.py:9:22
  %2 : int = prim::Constant[value=4]() # test.py:9:25
  %20 : int = prim::Constant[value=10]() # test.py:11:16
  %23 : float = prim::Constant[value=1]() # test.py:12:23
  %4 : int[] = prim::ListConstruct(%1, %2)
  %rv.1 : Tensor = aten::zeros(%4, %6, %6, %10, %12) # test.py:9:10
  %rv : Tensor = prim::Loop(%len.1, %17, %rv.1) # test.py:10:5
    block0(%i.1 : int, %rv.14 : Tensor):
      %21 : bool = aten::lt(%i.1, %20) # test.py:11:12
      %rv.13 : Tensor = prim::If(%21) # test.py:11:9
        block0():
          %rv.3 : Tensor = aten::sub(%rv.14, %23, %24) # test.py:12:18
          -> (%rv.3)
        block1():
          %rv.6 : Tensor = aten::add(%rv.14, %23, %24) # test.py:14:18
          -> (%rv.6)
      -> (%17, %rv.13)
  return (%rv)

以指令 %rv.1 : Tensor = aten::zeros(%4, %6, %6, %10, %12) # test.py:9:10 为例。

• %rv.1 : Tensor 表示我们将输出分配给名为 rv.1 的（唯一）值，该值的类型为 Tensor，并且我们不知道它的具体形状。

• aten::zeros 是运算符（相当于 torch.zeros），输入列表 (%4, %6, %6, %10, %12) 指定应将范围内的哪些值作为输入传递。可以在 内置函数 中找到内置函数（如 aten::zeros）的模式。

• # test.py:9:10 是生成此指令的原始源文件中的位置。在本例中，它是一个名为 test.py 的文件，位于第 9 行，字符 10 处。

请注意，运算符还可以关联 块，即 prim::Loop 和 prim::If 运算符。在图打印输出中，这些运算符的格式反映了它们等效的源代码形式，以便于调试。

可以按所示方式检查图，以确认 ScriptModule 描述的计算是正确的，如下面描述的自动和手动方式。

跟踪器

跟踪边缘情况

存在一些边缘情况，其中给定 Python 函数/模块的跟踪不能代表底层代码。这些情况可能包括

• 跟踪依赖于输入（例如张量形状）的控制流

• 跟踪张量视图的原地操作（例如赋值左侧的索引）

请注意，将来实际上可以跟踪这些情况。

自动跟踪检查

自动捕获跟踪中许多错误的一种方法是在 torch.jit.trace() API 上使用 check_inputs。check_inputs 接受将用于重新跟踪计算并验证结果的输入元组列表。例如

def loop_in_traced_fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

traced = torch.jit.trace(loop_in_traced_fn, inputs, check_inputs=check_inputs)

为我们提供以下诊断信息

ERROR: Graphs differed across invocations!
Graph diff:

            graph(%x : Tensor) {
            %1 : int = prim::Constant[value=0]()
            %2 : int = prim::Constant[value=0]()
            %result.1 : Tensor = aten::select(%x, %1, %2)
            %4 : int = prim::Constant[value=0]()
            %5 : int = prim::Constant[value=0]()
            %6 : Tensor = aten::select(%x, %4, %5)
            %result.2 : Tensor = aten::mul(%result.1, %6)
            %8 : int = prim::Constant[value=0]()
            %9 : int = prim::Constant[value=1]()
            %10 : Tensor = aten::select(%x, %8, %9)
        -   %result : Tensor = aten::mul(%result.2, %10)
        +   %result.3 : Tensor = aten::mul(%result.2, %10)
        ?          ++
            %12 : int = prim::Constant[value=0]()
            %13 : int = prim::Constant[value=2]()
            %14 : Tensor = aten::select(%x, %12, %13)
        +   %result : Tensor = aten::mul(%result.3, %14)
        +   %16 : int = prim::Constant[value=0]()
        +   %17 : int = prim::Constant[value=3]()
        +   %18 : Tensor = aten::select(%x, %16, %17)
        -   %15 : Tensor = aten::mul(%result, %14)
        ?     ^                                 ^
        +   %19 : Tensor = aten::mul(%result, %18)
        ?     ^                                 ^
        -   return (%15);
        ?             ^
        +   return (%19);
        ?             ^
            }

此消息向我们表明，在我们第一次跟踪它和我们使用 check_inputs 跟踪它时，计算有所不同。实际上，loop_in_traced_fn 主体内的循环取决于输入 x 的形状，因此当我们尝试另一个形状不同的 x 时，跟踪会有所不同。

在这种情况下，可以使用 torch.jit.script() 来捕获这种依赖于数据的控制流

def fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

scripted_fn = torch.jit.script(fn)
print(scripted_fn.graph)
#print(str(scripted_fn.graph).strip())

for input_tuple in [inputs] + check_inputs:
    torch.testing.assert_close(fn(*input_tuple), scripted_fn(*input_tuple))

这会产生

graph(%x : Tensor) {
    %5 : bool = prim::Constant[value=1]()
    %1 : int = prim::Constant[value=0]()
    %result.1 : Tensor = aten::select(%x, %1, %1)
    %4 : int = aten::size(%x, %1)
    %result : Tensor = prim::Loop(%4, %5, %result.1)
    block0(%i : int, %7 : Tensor) {
        %10 : Tensor = aten::select(%x, %1, %i)
        %result.2 : Tensor = aten::mul(%7, %10)
        -> (%5, %result.2)
    }
    return (%result);
}

跟踪器警告

跟踪器会针对跟踪计算中的几种有问题模式生成警告。例如，获取包含对张量切片（视图）进行原地赋值的函数的跟踪

def fill_row_zero(x):
    x[0] = torch.rand(*x.shape[1:2])
    return x

traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
print(traced.graph)

产生几个警告和一个简单返回输入的图

fill_row_zero.py:4: TracerWarning: There are 2 live references to the data region being modified when tracing in-place operator copy_ (possibly due to an assignment). This might cause the trace to be incorrect, because all other views that also reference this data will not reflect this change in the trace! On the other hand, if all other views use the same memory chunk, but are disjoint (e.g. are outputs of torch.split), this might still be safe.
    x[0] = torch.rand(*x.shape[1:2])
fill_row_zero.py:6: TracerWarning: Output nr 1. of the traced function does not match the corresponding output of the Python function. Detailed error:
Not within tolerance rtol=1e-05 atol=1e-05 at input[0, 1] (0.09115803241729736 vs. 0.6782537698745728) and 3 other locations (33.00%)
    traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
graph(%0 : Float(3, 4)) {
    return (%0);
}

我们可以通过修改代码来解决这个问题，不要使用原地更新，而是使用 torch.cat 在异地构建结果张量。

def fill_row_zero(x):
    x = torch.cat((torch.rand(1, *x.shape[1:2]), x[1:2]), dim=0)
    return x

traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
print(traced.graph)

常见问题解答

问：我想在 GPU 上训练模型并在 CPU 上进行推理。最佳实践是什么？

首先将模型从 GPU 转换为 CPU，然后保存它，如下所示

cpu_model = gpu_model.cpu()
sample_input_cpu = sample_input_gpu.cpu()
traced_cpu = torch.jit.trace(cpu_model, sample_input_cpu)
torch.jit.save(traced_cpu, "cpu.pt")

traced_gpu = torch.jit.trace(gpu_model, sample_input_gpu)
torch.jit.save(traced_gpu, "gpu.pt")

# ... later, when using the model:

if use_gpu:
  model = torch.jit.load("gpu.pt")
else:
  model = torch.jit.load("cpu.pt")

model(input)

建议这样做是因为跟踪器可能会见证在特定设备上创建张量，因此转换已加载的模型可能会产生意外的影响。在保存模型之前转换模型可确保跟踪器具有正确的设备信息。

问：如何在 ScriptModule 上存储属性？

假设我们有一个如下模型

import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.x = 2

    def forward(self):
        return self.x

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

如果实例化 Model，将导致编译错误，因为编译器不知道 x。有 4 种方法可以将 ScriptModule 上的属性告知编译器

1. nn.Parameter - 包装在 nn.Parameter 中的值将像在 nn.Module 上一样工作

2. register_buffer - 包装在 register_buffer 中的值将像在 nn.Module 上一样工作。这等效于 Tensor 类型的属性（参见 4）。

3. 常量 - 将类成员注释为 Final（或在类定义级别将其添加到名为 __constants__ 的列表中）会将包含的名称标记为常量。常量直接保存在模型的代码中。有关详细信息，请参阅内置常量。

4. 属性 - 可以将 支持的类型 的值添加为可变属性。大多数类型可以推断出来，但有些类型可能需要指定，有关详细信息，请参阅模块属性。

问：我想跟踪模块的方法，但我总是收到此错误

运行时错误： 无法 将 需要 梯度 的 张量 插入 为 常量。 请考虑 将其 设为 参数 或 输入， 或 分离 梯度

此错误通常意味着您正在跟踪的方法使用模块的参数，并且您传递的是模块的方法而不是模块实例（例如 my_module_instance.forward 与 my_module_instance）。

• 使用模块的方法调用 trace 会将模块参数（可能需要梯度）捕获为**常量**。

• 另一方面，使用模块实例（例如 my_module）调用 trace 会创建一个新模块并将参数正确复制到新模块中，以便它们可以在需要时累积梯度。

要跟踪模块上的特定方法，请参阅 torch.jit.trace_module

已知问题

如果您在 TorchScript 中使用 Sequential，则某些 Sequential 子模块的输入可能会被错误地推断为 Tensor，即使它们已被注释为其他类型。规范的解决方案是继承 nn.Sequential 并使用正确键入的输入重新声明 forward。

附录

迁移到 PyTorch 1.2 递归脚本 API

本节详细介绍了 PyTorch 1.2 中对 TorchScript 的更改。如果您是 TorchScript 的新手，可以跳过本节。PyTorch 1.2 对 TorchScript API 有两个主要更改。

1. torch.jit.script 现在将尝试递归编译它遇到的函数、方法和类。一旦您调用 torch.jit.script，编译就是“选择退出”而不是“选择加入”。

2. torch.jit.script(nn_module_instance) 现在是创建 ScriptModule 的首选方法，而不是继承 torch.jit.ScriptModule。这些更改结合起来提供了一个更简单、更易于使用的 API，用于将您的 nn.Module 转换为 ScriptModule，准备在非 Python 环境中进行优化和执行。

新的用法如下所示

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

my_model = Model()
my_scripted_model = torch.jit.script(my_model)

• 默认情况下，模块的 forward 会被编译。从 forward 调用的方法会按照它们在 forward 中的使用顺序延迟编译。

• 要编译除 forward 之外未从 forward 调用的方法，请添加 @torch.jit.export。

• 要阻止编译器编译方法，请添加 @torch.jit.ignore 或 @torch.jit.unused。 @ignore 将

• 方法保留为对 python 的调用，而 @unused 会将其替换为异常。 @ignored 无法导出； @unused 可以导出。

• 大多数属性类型都可以推断出来，因此不需要 torch.jit.Attribute。对于空容器类型，请使用 PEP 526 样式 的类注释来注释它们的类型。

• 可以使用 Final 类注释来标记常量，而不是将成员的名称添加到 __constants__ 中。

• 可以使用 Python 3 类型提示来代替 torch.jit.annotate

由于这些变化，以下项目被认为已被弃用，并且不应出现在新代码中

• @torch.jit.script_method 装饰器

• 继承自 torch.jit.ScriptModule 的类

• torch.jit.Attribute 包装类

• __constants__ 数组

• torch.jit.annotate 函数

模块

警告

@torch.jit.ignore 注解的行为在 PyTorch 1.2 中发生了变化。在 PyTorch 1.2 之前，@ignore 装饰器用于使函数或方法可以从导出的代码中调用。要恢复此功能，请使用 @torch.jit.unused()。@torch.jit.ignore 现在等效于 @torch.jit.ignore(drop=False)。有关详细信息，请参阅 @torch.jit.ignore 和 @torch.jit.unused。

当传递给 torch.jit.script 函数时，torch.nn.Module 的数据会被复制到 ScriptModule，并且 TorchScript 编译器会编译该模块。默认情况下，模块的 forward 会被编译。从 forward 调用的方法会按照它们在 forward 中使用的顺序延迟编译，以及任何 @torch.jit.export 方法。

torch.jit.export(fn)

此装饰器指示 nn.Module 上的方法用作 ScriptModule 的入口点，并且应该进行编译。

forward 隐式地被假定为入口点，因此它不需要此装饰器。从 forward 调用的函数和方法会在编译器看到它们时进行编译，因此它们也不需要此装饰器。

示例（在方法上使用 @torch.jit.export）

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def implicitly_compiled_method(self, x):
        return x + 99

    # `forward` is implicitly decorated with `@torch.jit.export`,
    # so adding it here would have no effect
    def forward(self, x):
        return x + 10

    @torch.jit.export
    def another_forward(self, x):
        # When the compiler sees this call, it will compile
        # `implicitly_compiled_method`
        return self.implicitly_compiled_method(x)

    def unused_method(self, x):
        return x - 20

# `m` will contain compiled methods:
#     `forward`
#     `another_forward`
#     `implicitly_compiled_method`
# `unused_method` will not be compiled since it was not called from
# any compiled methods and wasn't decorated with `@torch.jit.export`
m = torch.jit.script(MyModule())

函数

函数没有太大变化，如果需要，可以用 @torch.jit.ignore 或 torch.jit.unused 装饰它们。

# Same behavior as pre-PyTorch 1.2
@torch.jit.script
def some_fn():
    return 2

# Marks a function as ignored, if nothing
# ever calls it then this has no effect
@torch.jit.ignore
def some_fn2():
    return 2

# As with ignore, if nothing calls it then it has no effect.
# If it is called in script it is replaced with an exception.
@torch.jit.unused
def some_fn3():
  import pdb; pdb.set_trace()
  return 4

# Doesn't do anything, this function is already
# the main entry point
@torch.jit.export
def some_fn4():
    return 2

TorchScript 类

警告

TorchScript 类支持是实验性的。目前它最适合简单的类似记录的类型（想想一个带有方法的 NamedTuple）。

用户定义的 TorchScript 类 中的所有内容默认都会导出，如果需要，可以用 @torch.jit.ignore 装饰函数。

属性

TorchScript 编译器需要知道“模块属性”的类型。大多数类型可以从成员的值推断出来。空列表和字典无法推断其类型，并且必须使用 PEP 526 风格 的类注解对其类型进行注解。如果无法推断类型并且没有明确注解，则它不会作为属性添加到生成的 ScriptModule 中

旧 API

from typing import Dict
import torch

class MyModule(torch.jit.ScriptModule):
    def __init__(self):
        super().__init__()
        self.my_dict = torch.jit.Attribute({}, Dict[str, int])
        self.my_int = torch.jit.Attribute(20, int)

m = MyModule()

新 API

from typing import Dict

class MyModule(torch.nn.Module):
    my_dict: Dict[str, int]

    def __init__(self):
        super().__init__()
        # This type cannot be inferred and must be specified
        self.my_dict = {}

        # The attribute type here is inferred to be `int`
        self.my_int = 20

    def forward(self):
        pass

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

常量

Final 类型构造函数可用于将成员标记为“常量”。如果成员未标记为常量，它们将作为属性复制到生成的 ScriptModule。如果已知值是固定的，则使用 Final 可以带来优化机会，并提供额外的类型安全性。

旧 API

class MyModule(torch.jit.ScriptModule):
    __constants__ = ['my_constant']

    def __init__(self):
        super().__init__()
        self.my_constant = 2

    def forward(self):
        pass
m = MyModule()

新 API

from typing import Final

class MyModule(torch.nn.Module):

    my_constant: Final[int]

    def __init__(self):
        super().__init__()
        self.my_constant = 2

    def forward(self):
        pass

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

变量

容器假定类型为 Tensor 并且是非可选的（有关更多信息，请参阅“默认类型”）。以前，torch.jit.annotate 用于告诉 TorchScript 编译器类型应该是什么。现在支持 Python 3 风格的类型提示。

import torch
from typing import Dict, Optional

@torch.jit.script
def make_dict(flag: bool):
    x: Dict[str, int] = {}
    x['hi'] = 2
    b: Optional[int] = None
    if flag:
        b = 2
    return x, b

融合后端

有几个融合后端可用于优化 TorchScript 执行。CPU 上的默认融合器是 NNC，它可以执行 CPU 和 GPU 的融合。GPU 上的默认融合器是 NVFuser，它支持更广泛的运算符，并且已证明生成的内核具有更高的吞吐量。有关使用和调试的更多详细信息，请参阅 NVFuser 文档。

参考文献

• Python 语言参考覆盖范围
• TorchScript 不支持的 PyTorch 构造