目录
Shortcuts

TorchScript

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

我们提供工具来逐步将模型从纯 Python 程序过渡到 TorchScript 程序，该程序可以独立于 Python 运行，例如在独立的 C++ 程序中。这使得可以在 PyTorch 中使用熟悉的 Python 工具训练模型，然后通过 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

我们将能够步入 @torch.jit.script 函数作为普通的 Python 函数。要禁用特定函数的 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。

请注意，运算符也可以具有关联的 blocks，即 prim::Loop 和 prim::If 运算符。在图打印输出中，这些运算符的格式设置为反映其等效的源代码形式，以方便轻松调试。

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

Tracer

追踪边缘情况

存在一些边缘情况，其中给定 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);
}

Tracer 警告

追踪器为追踪计算中的几种有问题模式生成警告。例如，获取包含张量的切片（视图）上的就地赋值的函数的追踪

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__ 的列表中）会将包含的名称标记为常量。常量直接保存在模型的代码中。有关详细信息，请参阅builtin-constants。

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

问：我想追踪模块的方法，但我一直收到此错误

RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient

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

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

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

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

已知问题

如果您将 Sequential 与 TorchScript 一起使用，即使某些 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 继承。这些更改结合在一起，为将您的 nn.Module 转换为 ScriptModule 提供了一个更简单、更易于使用的 API，以便在非 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 调用的方法以及任何 @torch.jit.export 方法都会按照它们在 forward 中使用的顺序进行延迟编译。

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 文档。

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得您的问题解答

查看资源

为了分析流量并优化您的体验，我们在本网站上使用 Cookie。通过点击或导航，您同意允许我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括关于可用控件：Cookie 政策。