torch.jit.trace#

torch.jit.trace(func, example_inputs=None, optimize=None, check_trace=True, check_inputs=None, check_tolerance=1e-05, strict=True, _force_outplace=False, _module_class=None, _compilation_unit=<torch.jit.CompilationUnit object>, example_kwarg_inputs=None, _store_inputs=True)[源]#

追踪一个函数并返回一个可执行的或 ScriptFunction，它将使用即时编译进行优化。

追踪非常适合仅作用于 Tensor 以及 Tensor 的列表、字典和元组的代码。

使用 torch.jit.trace 和 torch.jit.trace_module，您可以将现有模块或 Python 函数转换为 TorchScript ScriptFunction 或 ScriptModule。您必须提供示例输入，我们将运行该函数，并记录在所有张量上执行的操作。

独立函数的追踪结果会生成 ScriptFunction。
nn.Module.forward 或 nn.Module 的追踪结果会生成 ScriptModule。

此模块还包含原始模块拥有的所有参数。

警告

追踪只正确记录那些不依赖于数据（例如，张量中的数据没有条件）且没有任何未跟踪的外部依赖项（例如，执行输入/输出或访问全局变量）的函数和模块。追踪只记录在给定张量上运行给定函数时执行的操作。因此，返回的 ScriptModule 将始终在任何输入上运行相同的追踪图。当您的模块预期根据输入和/或模块状态运行不同的操作集时，这会产生一些重要的影响。例如，

追踪不会记录任何控制流，如 if 语句或循环。当这种控制流在您的模块中是常量时，这很好，并且它通常会内联控制流决策。但有时控制流实际上是模型本身的一部分。例如，循环网络是对输入序列（可能动态的）长度的循环。
在返回的 ScriptModule 中，在 training 和 eval 模式下具有不同行为的操作将始终表现为在追踪时的模式，无论 ScriptModule 当前处于哪种模式。

在这些情况下，追踪可能不适用，脚本化 是一个更好的选择。如果您追踪此类模型，您可能会在后续模型调用中静默地得到不正确的结果。当执行可能导致生成不正确追踪的操作时，追踪器将尝试发出警告。

参数

func (callable 或 torch.nn.Module) – 一个将使用 example_inputs 运行的 Python 函数或 torch.nn.Module。func 的参数和返回值必须是张量或（可能是嵌套的）包含张量的元组。当将模块传递给 torch.jit.trace 时，只会运行和追踪 forward 方法（有关详细信息，请参阅 torch.jit.trace）。

关键字参数

example_inputs (tuple 或 torch.Tensor 或 None, 可选) – 在跟踪时将传递给函数的示例输入元组。默认值：None。必须指定此参数或example_kwarg_inputs。假定跟踪的操作支持这些类型和形状，则生成的跟踪可以使用不同类型和形状的输入运行。example_inputs 也可以是单个张量，在这种情况下，它会自动包装在元组中。当值为None时，应指定example_kwarg_inputs。
check_trace (bool, 可选) – 检查相同输入通过跟踪代码是否生成相同输出。默认值：True。如果您网络包含非确定性操作，或者您确定网络即使检查失败也是正确的，则可能需要禁用此功能。
check_inputs (list of tuples, 可选) – 应使用的一组输入参数的元组列表，用于检查跟踪是否与预期一致。每个元组等同于一组将在example_inputs中指定的输入参数。为获得最佳结果，请传入一组代表您期望网络看到的输入形状和类型空间的检查输入。如果未指定，则使用原始的example_inputs进行检查。
check_tolerance (float, 可选) – 在检查过程中使用的浮点比较容差。当由于已知原因（例如操作融合）导致结果在数值上发散时，可以使用此参数来放宽检查器的严格性。
strict (bool, 可选) – 是否以严格模式运行跟踪器（默认：True）。仅当您希望跟踪器记录您的可变容器类型（目前为list/dict），并且您确定问题中使用的容器是constant结构并且不作为控制流（if, for）条件使用时，才将其关闭。
example_kwarg_inputs (dict, 可选) – 此参数是示例输入的关键字参数包，在跟踪时将传递给函数。默认值：None。必须指定此参数或example_inputs。该字典将按跟踪函数的参数名称进行解包。如果字典的键与跟踪函数的参数名称不匹配，则会引发运行时异常。

返回

如果func是nn.Module或nn.Module的forward，则trace返回一个ScriptModule对象，其中包含一个带有跟踪代码的forward方法。返回的ScriptModule将拥有与原始nn.Module相同的子模块和参数集。如果func是独立函数，则trace返回ScriptFunction。

示例（跟踪函数）

import torch

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

# Run `foo` with the provided inputs and record the tensor operations
traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

# `traced_foo` can now be run with the TorchScript interpreter or saved
# and loaded in a Python-free environment

示例（跟踪现有模块）

import torch
import torch.nn as nn


class Net(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = nn.Conv2d(1, 1, 3)

    def forward(self, x):
        return self.conv(x)


n = Net()
example_weight = torch.rand(1, 1, 3, 3)
example_forward_input = torch.rand(1, 1, 3, 3)

# Trace a specific method and construct `ScriptModule` with
# a single `forward` method
module = torch.jit.trace(n.forward, example_forward_input)

# Trace a module (implicitly traces `forward`) and construct a
# `ScriptModule` with a single `forward` method
module = torch.jit.trace(n, example_forward_input)

torch.jit.trace#

文档

教程

资源