快捷方式

torch.jit.trace_module

torch.jit.trace_module(mod, inputs, 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_inputs_is_kwarg=False, _store_inputs=True)[源][源]

追踪一个模块并返回一个可执行的 ScriptModule,该模块将使用即时编译进行优化。

当一个模块被传递给 torch.jit.trace 时,只会运行和追踪 forward 方法。使用 trace_module,你可以指定一个方法名到示例输入的字典来追踪(见下面的 inputs 参数)。

有关追踪的更多信息,请参阅 torch.jit.trace

参数
  • mod (torch.nn.Module) – 一个包含 inputs 中指定方法名称的 torch.nn.Module。给定的方法将作为单个 ScriptModule 的一部分进行编译。

  • inputs (dict) – 一个字典,包含按 mod 中方法名称索引的示例输入。这些输入在追踪时将传递给与输入键对应的方法。例如:{ 'forward' : example_forward_input, 'method2': example_method2_input}

关键字参数
  • check_trace (bool, 可选) – 检查通过追踪代码运行的相同输入是否产生相同的输出。默认值:True。如果您的网络包含非确定性操作,或者您确定网络正确尽管检查器失败,您可能希望禁用此选项。

  • check_inputs (list of dicts, optional) – 一个包含输入参数字典的列表,用于对照预期结果检查追踪。每个元组等同于在 inputs 中指定的一组输入参数。为了获得最佳结果,请传入一组具有代表性的形状和类型输入的检查输入,这些输入代表您期望网络看到的输入空间。如果未指定,则使用原始 inputs 进行检查。

  • check_tolerance (float, optional) – 在检查过程中使用的浮点比较容差。这可用于在结果由于已知原因(例如运算符融合)出现数值差异时放宽检查器严格性。

  • example_inputs_is_kwarg (bool, 可选) – 此参数指示示例输入是否为关键字参数包。默认值:False

返回

一个 ScriptModule 对象,其中包含一个具有追踪代码的 forward 方法。当 func 是一个 torch.nn.Module 时,返回的 ScriptModule 将具有与 func 相同的子模块和参数集。

示例 (追踪具有多个方法的模块)

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)

    def weighted_kernel_sum(self, weight):
        return weight * self.conv.weight


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)

# Trace specific methods on a module (specified in `inputs`), constructs
# a `ScriptModule` with `forward` and `weighted_kernel_sum` methods
inputs = {"forward": example_forward_input, "weighted_kernel_sum": example_weight}
module = torch.jit.trace_module(n, inputs)

文档

查阅 PyTorch 全面的开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源