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>) [source]

Отследить модуль и вернуть исполняемый ScriptModule, который будет оптимизирован с помощью компиляции в реальном времени. Когда модуль передаётся в torch.jit.trace, выполняется и отслеживается только метод forward. С помощью trace_module, можно указать словарь имён методов и входных данных для отслеживания (см. аргумент inputs) ниже.

См. torch.jit.trace для получения дополнительной информации об отслеживании.

Параметры:

• mod (torch.nn.Module) – Модуль torch.nn.Module содержащий методы, имена которых указаны в inputs. Указанные методы будут скомпилированы как часть одного ScriptModule.
• inputs (dict) – Словарь, содержащий примерные входные данные, индексированные по именам методов в mod. Входные данные будут переданы методам, имена которых соответствуют ключам входных данных во время отслеживания. { 'forward' : example_forward_input, 'method2': example_method2_input}

Ключевые аргументы:

• check_trace (bool, optional) – Проверка, что одни и те же входные данные, проходящие через отслеженный код, дают одни и те же выходные данные. По умолчанию: True. Возможно, вы захотите отключить это, если, например, ваша сеть содержит недетерминированные операции или если вы уверены, что сеть правильная, несмотря на ошибку проверки.
• check_inputs (список словарей, необязательно) – Список словарей входных аргументов, которые должны быть использованы для проверки отслеживания по сравнению с ожидаемым результатом. Каждый кортеж эквивалентен набору входных аргументов, которые были бы указаны в inputs. Для достижения наилучших результатов передайте набор проверочных входных данных, репрезентативных для пространства форм и типов входных данных, которые, как ожидается, увидит сеть. Если не указано, исходные inputs используются для проверки
• check_tolerance (float, необязательно) – Допуск для сравнения чисел с плавающей запятой, используемый в процедуре проверки. Это может использоваться для ослабления строгости проверки в случае, если результаты расходятся численно по известной причине, например, из-за слияния операций.

Возвращает:

Объект ScriptModule с единственным методом forward, содержащим отслеженный код. Если func является torch.nn.Module, возвращаемый ScriptModule будет иметь тот же набор подмодулей и параметров, что и func.

Пример (отслеживание модуля с несколькими методами):

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)