Python PyTorch trace_module用法及代码示例

本文简要介绍python语言中 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>)

参数：

• mod(torch.nn.Module) -torch.nn.Module 包含名称在 inputs 中指定的方法。给定的方法将被编译为单个 ScriptModule 的一部分。

• inputs(dict) -包含由 mod 中的方法名称索引的示例输入的字典。输入将在跟踪时传递给名称对应于输入键的方法。 { 'forward' : example_forward_input, 'method2': example_method2_input}

关键字参数：

• check_trace(bool， 可选的) -检查通过跟踪代码运行的相同输入是否产生相同的输出。默认值：True。例如，如果您的网络包含非确定性操作，或者您确定网络是正确的(尽管检查器失败)，您可能希望禁用此函数。

• check_inputs(字典列表,可选的) -输入参数的字典列表，应用于检查跟踪是否符合预期。每个元组等效于将在 inputs 中指定的一组输入参数。为了获得最佳结果，请传入一组检查输入，这些输入代表您希望网络看到的形状和输入类型的空间。如果未指定，则使用原始inputs 进行检查

• check_tolerance(float,可选的) -在检查程序中使用的浮点比较容差。如果由于已知原因(例如运算符融合)导致结果在数值上出现分歧，这可以用来放宽检查器的严格性。

返回：

一个 ScriptModule 对象，带有一个包含跟踪代码的 forward 方法。当 func 是 torch.nn.Module 时，返回的 ScriptModule 将具有与 func 相同的一组 sub-modules 和参数。

跟踪一个模块并返回一个可执行文件 ScriptModule ，它将使用just-in-time 编译进行优化。将模块传递给 torch.jit.trace 时，仅运行和跟踪 forward 方法。使用 trace_module ，您可以将方法名称字典指定为示例输入以跟踪以下参数(请参阅 inputs )。

有关跟踪的更多信息，请参阅 torch.jit.trace 。

示例(使用多种方法跟踪模块)：

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__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)