torch.onnx

Пример: AlexNet из PyTorch в ONNX

Вот простой скрипт, который экспортирует предобученную AlexNet в файл ONNX с именем alexnet.onnx. Вызов torch.onnx.export выполняет модель один раз, чтобы отследить ее выполнение, а затем экспортирует отслеженную модель в указанный файл:

import torch
import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device="cuda")
model = torchvision.models.alexnet(pretrained=True).cuda()

# Providing input and output names sets the display names for values
# within the model's graph. Setting these does not change the semantics
# of the graph; it is only for readability.
#
# The inputs to the network consist of the flat list of inputs (i.e.
# the values you would pass to the forward() method) followed by the
# flat list of parameters. You can partially specify names, i.e. provide
# a list here shorter than the number of inputs to the model, and we will
# only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

Полученный файл alexnet.onnx содержит двоичный протокольный буфер, который содержит как структуру сети, так и параметры экспортированной модели (в данном случае AlexNet). Аргумент verbose=True заставляет экспортер вывести удобочитаемое представление модели:

# These are the inputs and parameters to the network, which have taken on
# the names we specified earlier.
graph(%actual_input_1 : Float(10, 3, 224, 224)
      %learned_0 : Float(64, 3, 11, 11)
      %learned_1 : Float(64)
      %learned_2 : Float(192, 64, 5, 5)
      %learned_3 : Float(192)
      # ---- omitted for brevity ----
      %learned_14 : Float(1000, 4096)
      %learned_15 : Float(1000)) {
  # Every statement consists of some output tensors (and their types),
  # the operator to be run (with its attributes, e.g., kernels, strides,
  # etc.), its input tensors (%actual_input_1, %learned_0, %learned_1)
  %17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
  %18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
  %19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
  # ---- omitted for brevity ----
  %29 : Float(10, 256, 6, 6) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%28), scope: AlexNet/Sequential[features]/MaxPool2d[12]
  # Dynamic means that the shape is not known. This may be because of a
  # limitation of our implementation (which we would like to fix in a
  # future release) or shapes which are truly dynamic.
  %30 : Dynamic = onnx::Shape(%29), scope: AlexNet
  %31 : Dynamic = onnx::Slice[axes=[0], ends=[1], starts=[0]](%30), scope: AlexNet
  %32 : Long() = onnx::Squeeze[axes=[0]](%31), scope: AlexNet
  %33 : Long() = onnx::Constant[value={9216}](), scope: AlexNet
  # ---- omitted for brevity ----
  %output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
  return (%output1);
}

Вы также можете проверить вывод с помощью библиотеки ONNX, которую можно установить, используя pip:

pip install onnx

Затем вы можете выполнить:

import onnx

# Load the ONNX model
model = onnx.load("alexnet.onnx")

# Check that the model is well formed
onnx.checker.check_model(model)

# Print a human readable representation of the graph
print(onnx.helper.printable_graph(model.graph))

Вы также можете запустить экспортированную модель с помощью одного из множества времени выполнения, поддерживающих ONNX. Например, после установки ONNX Runtime, вы можете загрузить и запустить модель:

import onnxruntime as ort

ort_session = ort.InferenceSession("alexnet.onnx")

outputs = ort_session.run(
    None,
    {"actual_input_1": np.random.randn(10, 3, 224, 224).astype(np.float32)},
)
print(outputs[0])

Вот более подробный учебник по экспорту модели и запуску с ONNX Runtime.

Трейсинг против скриптинга

Внутренне, torch.onnx.export() требует torch.jit.ScriptModule, а не torch.nn.Module. Если переданная модель не является уже ScriptModule, export() будет использовать трейсинг для ее преобразования:

• Трейсинг: Если torch.onnx.export() вызывается с модулем, который еще не является ScriptModule, он сначала выполняет эквивалент torch.jit.trace(), который выполняет модель один раз с заданным args и записывает все операции, происходящие во время этого выполнения. Это означает, что если ваша модель динамическая, например, меняет поведение в зависимости от входных данных, экспортированная модель не захватывает это динамическое поведение. Рекомендуется изучить экспортированную модель и убедиться, что операторы выглядят разумно. Трейсинг будет развертывать циклы и условные операторы if, экспортируя статическую граф, который точно такой же, как отслеженный запуск. Если вы хотите экспортировать свою модель с динамическим потоком управления, вам нужно использовать скриптинг.

• Скриптинг: Компиляция модели через скриптинг сохраняет динамический поток управления и допустима для входных данных разных размеров. Чтобы использовать скриптинг:

  • Используйте torch.jit.script() для создания ScriptModule.
  • Вызовите torch.onnx.export() с ScriptModule в качестве модели. args все еще требуются, но они будут использоваться только внутри для создания примеров вывода, чтобы можно было захватить типы и формы вывода. Никакой трейсинг не будет выполнен.

См. Введение в TorchScript и TorchScript для получения дополнительных сведений, в том числе о том, как объединить трейсинг и скриптинг, чтобы удовлетворить конкретным требованиям различных моделей.

Избегание ошибок

# Bad! Will be replaced with constants during tracing.
x, y = np.random.rand(1, 2), np.random.rand(1, 2)
np.concatenate((x, y), axis=1)

# Good! Tensor operations will be captured during tracing.
x, y = torch.randn(1, 2), torch.randn(1, 2)
torch.cat((x, y), dim=1)

# Bad! y.item() will be replaced with a constant during tracing.
def forward(self, x, y):
    return x.reshape(y.item(), -1)

# Good! y will be preserved as a variable during tracing.
def forward(self, x, y):
    return x.reshape(y, -1)

class Model(torch.nn.Module):
  def forward(self, states):
      batch_size, seq_length = states.shape[:2]
      real_seq_length = seq_length
      real_seq_length += 2
      return real_seq_length + seq_length

real_seq_length = real_seq_length + 2

Ограничения

# Tensor indices that includes negative values.
data[torch.tensor([[1, 2], [2, -3]]), torch.tensor([-2, 3])]
# Workarounds: use positive index values.

# Multiple tensor indices if any has rank >= 2
data[torch.tensor([[1, 2], [2, 3]]), torch.tensor([2, 3])] = new_data
# Workarounds: use single tensor index with rank >= 2,
#              or multiple consecutive tensor indices with rank == 1.

# Multiple tensor indices that are not consecutive
data[torch.tensor([2, 3]), :, torch.tensor([1, 2])] = new_data
# Workarounds: transpose `data` such that tensor indices are consecutive.

# Tensor indices that includes negative values.
data[torch.tensor([1, -2]), torch.tensor([-2, 3])] = new_data
# Workarounds: use positive index values.

# Implicit broadcasting required for new_data.
data[torch.tensor([[0, 2], [1, 1]]), 1:3] = new_data
# Workarounds: expand new_data explicitly.
# Example:
#   data shape: [3, 4, 5]
#   new_data shape: [5]
#   expected new_data shape after broadcasting: [2, 2, 2, 5]

Добавление поддержки операторов

При экспорте модели, содержащей неподдерживаемые операторы, вы увидите сообщение об ошибке, подобное:

RuntimeError: ONNX export failed: Couldn't export operator foo

В таких случаях вы можете предпринять несколько действий:

1. Изменить модель так, чтобы она не использовала этот оператор.
2. Создать символическую функцию для преобразования оператора и зарегистрировать ее как пользовательскую символическую функцию.
3. Внести вклад в PyTorch, чтобы добавить ту же символическую функцию в torch.onnx .

Если вы решили реализовать символическую функцию (мы надеемся, что вы внесете свой вклад в PyTorch!), вот как вы можете начать:

def foo(
  g,
  input_0: torch._C.Value,
  input_1: torch._C.Value) -> Union[None, torch._C.Value, List[torch._C.Value]]:
  """
  Adds the ONNX operations representing this PyTorch function by updating the
  graph g with `g.op()` calls.

  Args:
    g (Graph): graph to write the ONNX representation into.
    input_0 (Value): value representing the variables which contain
        the first input for this operator.
    input_1 (Value): value representing the variables which contain
        the second input for this operator.

  Returns:
    A Value or List of Values specifying the ONNX nodes that compute something
    equivalent to the original PyTorch operator with the given inputs.

    None if it cannot be converted to ONNX.
  """
  ...

print(
    torch.jit.trace(
        torch.nn.ELU(), # module
        torch.ones(1)   # example input
    ).graph
)

graph(%self : __torch__.torch.nn.modules.activation.___torch_mangle_0.ELU,
      %input : Float(1, strides=[1], requires_grad=0, device=cpu)):
  %4 : float = prim::Constant[value=1.]()
  %5 : int = prim::Constant[value=1]()
  %6 : int = prim::Constant[value=1]()
  %7 : Float(1, strides=[1], requires_grad=0, device=cpu) = aten::elu(%input, %4, %5, %6)
  return (%7)

def elu(input: Tensor, alpha: float = ..., inplace: bool = ...) -> Tensor: ...

def elu(g, input: torch.Value, alpha: torch.Value, inplace: bool = False):
    return g.op("Elu", input, alpha_f=alpha)

class MyRelu(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def symbolic(g: torch.Graph, input: torch.Value) -> torch.Value:
        return g.op("Clip", input, g.op("Constant", value_t=torch.tensor(0, dtype=torch.float)))

class MyLogExp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input: torch.Tensor) -> torch.Tensor:
        ctx.save_for_backward(input)
        h = input.exp()
        return h.log().log()

graph(%input : Float(1, strides=[1], requires_grad=0, device=cpu)):
    %1 : float = onnx::Exp[](%input)
    %2 : float = onnx::Log[](%1)
    %3 : float = onnx::Log[](%2)
    return (%3)

from torch.onnx import symbolic_helper


# Define custom symbolic function
@symbolic_helper.parse_args("v", "v", "f", "i")
def symbolic_foo_forward(g, input1, input2, attr1, attr2):
    return g.op("custom_domain::Foo", input1, input2, attr1_f=attr1, attr2_i=attr2)


# Register custom symbolic function
torch.onnx.register_custom_op_symbolic("custom_ops::foo_forward", symbolic_foo_forward, 9)


class FooModel(torch.nn.Module):
    def __init__(self, attr1, attr2):
        super().__init__()
        self.attr1 = attr1
        self.attr2 = attr2

    def forward(self, input1, input2):
        # Calling custom op
        return torch.ops.custom_ops.foo_forward(input1, input2, self.attr1, self.attr2)


model = FooModel(attr1, attr2)
torch.onnx.export(
    model,
    (example_input1, example_input1),
    "model.onnx",
    # only needed if you want to specify an opset version > 1.
    custom_opsets={"custom_domain": 2}
)

# prepare model, args, opset_version
...

torch_script_graph, unconvertible_ops = torch.onnx.utils.unconvertible_ops(
    model, args, opset_version=opset_version
)

print(set(unconvertible_ops))

Часто задаваемые вопросы

В: Я экспортировал свою модель LSTM, но размер её входных данных, кажется, фиксирован?

Трейсер записывает размер входных данных примера. Если модель должна принимать входные данные с динамическими размерами, установите dynamic_axes при вызове torch.onnx.export().

В: Как экспортировать модели, содержащие циклы?

См. Трейсинг против скриптинга.

В: Как экспортировать модели с входными данными примитивных типов (например, int, float)?

Поддержка входных данных примитивных числовых типов была добавлена в PyTorch 1.9. Однако экспортер не поддерживает модели со строковыми входными данными.

В: Поддерживает ли ONNX неявное приведение типов скаляров?

Стандарт ONNX этого не поддерживает, но экспортер попытается справиться с этим. Скаляры экспортируются как тензоры-константы. Экспортер определит правильный тип данных для скаляров. В редких случаях, когда он не сможет сделать это, вам необходимо вручную указать тип данных, например, с помощью dtype=torch.float32. Если вы увидите какие-либо ошибки, пожалуйста, [создайте запрос в GitHub](https://github.com/pytorch/pytorch/issues).

В: Могут ли списки тензоров экспортироваться в ONNX?

Да, для opset_version >= 11, так как ONNX представил тип Sequence в опсет 11.

Содействие/разработка

Документация для разработчиков.

Функции

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=<TrainingMode.EVAL: 0>, input_names=None, output_names=None, operator_export_type=<OperatorExportTypes.ONNX: 0>, opset_version=None, do_constant_folding=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, export_modules_as_functions=False) [source]

Экспортирует модель в формате ONNX.

Если model не является torch.jit.ScriptModule и не torch.jit.ScriptFunction, это выполняет model один раз для преобразования в граф TorchScript для экспорта (эквивалентно torch.jit.trace()). Таким образом, это имеет те же ограниченные возможности поддержки динамического управления потоком, что и torch.jit.trace().

Параметры:

• model (torch.nn.Module, torch.jit.ScriptModule или torch.jit.ScriptFunction) – модель, подлежащая экспорту.

• args (кортеж или torch.Tensor) –

  Аргументы могут быть структурированы следующим образом:

  1. ТОЛЬКО КОРТЕЖ АРГУМЕНТОВ:

     args = (x, y, z)
     

  Кортеж должен содержать входные данные модели таким образом, что model(*args) является допустимым вызовом модели. Любые аргументы, не являющиеся тензорами, будут жёстко закодированы в экспортируемую модель; любые тензорные аргументы станут входными данными экспортируемой модели в порядке их появления в кортеже.

  2. ТЕНЗОР:

     args = torch.Tensor([1])
     

  Это эквивалентно кортежу с одним элементом — этим тензором.

  3. КОРТЕЖ АРГУМЕНТОВ, ЗАВЕРШАЮЩИЙСЯ СЛОВАРЕМ ИМЕНЕННЫХ АРГУМЕНТОВ:

     args = (
         x,
         {
             "y": input_y,
             "z": input_z
         }
     )
     

  Все элементы кортежа, кроме последнего, будут переданы как позиционные аргументы, а именованные аргументы будут установлены из последнего элемента. Если именованный аргумент отсутствует в словаре, ему будет присвоено значение по умолчанию или None, если значение по умолчанию не задано.

  Примечание

  Если словарь является последним элементом кортежа args, он будет интерпретирован как содержащий именованные аргументы. Чтобы передать словарь как последний позиционный аргумент, передайте пустой словарь в качестве последнего элемента кортежа args. Например, вместо:

  torch.onnx.export(
      model,
      (
          x,
          # WRONG: will be interpreted as named arguments
          {y: z}
      ),
      "test.onnx.pb"
  )
  

  Напишите:

  torch.onnx.export(
      model,
      (
          x,
          {y: z},
          {}
      ),
      "test.onnx.pb"
  )
  

• f (Объединение[строка, BytesIO]) – объект типа "поток" (такой, что f.fileno() возвращает дескриптор файла) или строка, содержащая имя файла. В этот файл будет записан двоичный протокол буфера.
• export_params (булево, значение по умолчанию True) – если True, все параметры будут экспортированы. Установите это значение в False, если вы хотите экспортировать не обученную модель. В этом случае экспортированная модель сначала примет все свои параметры в качестве аргументов, в порядке, указанном model.state_dict().values()
• verbose (булево, значение по умолчанию False) – если True, печатает описание экспортируемой модели в стандартный вывод. Кроме того, конечная графика ONNX будет содержать поле doc_string` из экспортируемой модели, которое указывает на места в исходном коде для model. Если True, ведение журнала экспортера ONNX будет включено.

• training (перечисление, значение по умолчанию TrainingMode.EVAL) –

  • TrainingMode.EVAL: экспорт модели в режиме вывода.

  • TrainingMode.PRESERVE: export the model in inference mode if model.training is

    False и в обучающем режиме, если model.training равно True.

  • TrainingMode.TRAINING: export the model in training mode. Disables optimizations

    что может помешать обучению.

• input_names (список строк, значение по умолчанию пустой список) – имена, присваиваемые входным узлам графа в порядке.
• output_names (список строк, значение по умолчанию пустой список) – имена, присваиваемые выходным узлам графа в порядке.

• operator_export_type (перечисление, значение по умолчанию OperatorExportTypes.ONNX) –

  • OperatorExportTypes.ONNX: Export all ops as regular ONNX ops

    (в стандартном домене opset).

  • OperatorExportTypes.ONNX_FALLTHROUGH: Try to convert all ops

    в стандартные операции ONNX в стандартном домене opset. Если это невозможно (например, потому что поддержка не была добавлена для преобразования определенной операции torch в ONNX), то операция экспортируется в домен настраиваемого opset без преобразования. Применимо к настраиваемым операциям, а также операциям ATen. Для того, чтобы экспортируемая модель была работоспособна, среда выполнения должна поддерживать эти нестандартные операции.

  • OperatorExportTypes.ONNX_ATEN: All ATen ops (in the TorchScript namespace “aten”)

    экспортируются как операции ATen (в домене opset «org.pytorch.aten»). ATen — встроенная библиотека тензоров PyTorch, поэтому это инструкция для среды выполнения использовать реализацию этих операций PyTorch.

    Предупреждение

    Модели, экспортированные таким образом, вероятно, будут работоспособны только в Caffe2.

    Это может быть полезно, если числовые различия в реализациях операторов вызывают большие различия в поведении между PyTorch и Caffe2 (что чаще встречается в необученных моделях).

  • OperatorExportTypes.ONNX_ATEN_FALLBACK: Try to export each ATen op

    (в пространстве имен TorchScript «aten») как обычная операция ONNX. Если это невозможно (например, потому что поддержка не была добавлена для преобразования определенной операции torch в ONNX), то операция экспортируется как операция ATen. См. документацию по OperatorExportTypes.ONNX_ATEN для контекста. Например:

    graph(%0 : Float):
    %3 : int = prim::Constant[value=0]()
    # conversion unsupported
    %4 : Float = aten::triu(%0, %3)
    # conversion supported
    %5 : Float = aten::mul(%4, %0)
    return (%5)
    

    Предполагая, что aten::triu не поддерживается в ONNX, это будет экспортировано как:

    graph(%0 : Float):
    %1 : Long() = onnx::Constant[value={0}]()
    # not converted
    %2 : Float = aten::ATen[operator="triu"](%0, %1)
    # converted
    %3 : Float = onnx::Mul(%2, %0)
    return (%3)
    

    Если PyTorch был скомпилирован с Caffe2 (т.е. с BUILD_CAFFE2=1), то будет включено поведение, специфичное для Caffe2, включая специальную поддержку операций, производимых модулями, описанными в Квантовании.

    Предупреждение

    Модели, экспортированные таким образом, вероятно, будут работоспособны только в Caffe2.

• opset_version (целое число, значение по умолчанию 14) – версия стандартного (ai.onnx) opset, на который ориентироваться. Должно быть >= 7 и <= 16.
• do_constant_folding (булево, значение по умолчанию True) – применить оптимизацию константного складывания. Константное складывание заменит некоторые операции, имеющие все константные входные данные, на вычисленные константные узлы.

• dynamic_axes (словарь[строка, словарь[целое число, строка]] или словарь[строка, список(целое число)], значение по умолчанию пустой словарь) –

  По умолчанию экспортируемая модель будет иметь формы всех тензоров входных и выходных данных, точно соответствующие тем, которые указаны в args. Чтобы указать оси тензоров как динамические (т.е. известные только во время выполнения), установите dynamic_axes в словарь со схемой:

  • KEY (str): an input or output name. Each name must also be provided in input_names or

    output_names.

  • ЗНАЧЕНИЕ (словарь или список): если словарь, ключи — индексы осей, значения — имена осей. Если список, каждый элемент — индекс оси.

  Например:

  class SumModule(torch.nn.Module):
      def forward(self, x):
          return torch.sum(x, dim=1)
  
  torch.onnx.export(
      SumModule(),
      (torch.ones(2, 2),),
      "onnx.pb",
      input_names=["x"],
      output_names=["sum"]
  )
  

  Производит:

  input {
    name: "x"
    ...
        shape {
          dim {
            dim_value: 2  # axis 0
          }
          dim {
            dim_value: 2  # axis 1
  ...
  output {
    name: "sum"
    ...
        shape {
          dim {
            dim_value: 2  # axis 0
  ...
  

  В то время как:

  torch.onnx.export(
      SumModule(),
      (torch.ones(2, 2),),
      "onnx.pb",
      input_names=["x"],
      output_names=["sum"],
      dynamic_axes={
          # dict value: manually named axes
          "x": {0: "my_custom_axis_name"},
          # list value: automatic names
          "sum": [0],
      }
  )
  

  Производит:

  input {
    name: "x"
    ...
        shape {
          dim {
            dim_param: "my_custom_axis_name"  # axis 0
          }
          dim {
            dim_value: 2  # axis 1
  ...
  output {
    name: "sum"
    ...
        shape {
          dim {
            dim_param: "sum_dynamic_axes_1"  # axis 0
  ...
  

• keep_initializers_as_inputs (булево, значение по умолчанию None) –

  Если True, все инициализаторы (как правило, соответствующие параметрам) в экспортируемой графике также будут добавлены в качестве входных данных в графу. Если False, инициализаторы не добавляются в качестве входных данных в граф, и в качестве входных данных добавляются только непараметрические входные данные. Это может позволить бэкендам/средам выполнения лучше оптимизировать (например, константное складывание).

  Если opset_version < 9, инициализаторы ДОЛЖНЫ быть частью входных данных графа, и этот аргумент будет проигнорирован, и поведение будет эквивалентно установлению этого аргумента в True.

  Если None, то поведение выбирается автоматически следующим образом:

  • If operator_export_type=OperatorExportTypes.ONNX, the behavior is equivalent

    эквивалентно установлению этого аргумента в False.

  • В противном случае поведение эквивалентно установлению этого аргумента в True.

• custom_opsets (словарь[строка, целое число], значение по умолчанию пустой словарь) –

  Словарь со схемой:

  • КЛЮЧ (строка): имя домена opset
  • ЗНАЧЕНИЕ (целое число): версия opset

  Если настраиваемый opset упоминается в model, но не указан в этом словаре, версия opset устанавливается в 1. Только имя и версия домена настраиваемого opset должны быть указаны с помощью этого аргумента.

• export_modules_as_functions (bool или множество типов python:type nn.Module, по умолчанию False) –

  Флаг для включения экспорта всех вызовов nn.Module forward в виде локальных функций в ONNX. Или множество для указания конкретных типов модулей для экспорта в виде локальных функций в ONNX. Эта функция требует opset_version >= 15, в противном случае экспорт завершится ошибкой. Это связано с тем, что opset_version < 15 подразумевает версию IR < 8, что означает отсутствие поддержки локальных функций. Переменные модулей будут экспортированы в качестве атрибутов функций. Существует две категории атрибутов функций.

  1. Аннотированные атрибуты: переменные класса, которые имеют аннотации типа через PEP 526, будут экспортированы как атрибуты. Аннотированные атрибуты не используются внутри подграфа локальной функции ONNX, потому что они не создаются трассировкой PyTorch JIT, но они могут использоваться потребителями, чтобы определить, нужно ли заменить функцию конкретным объединённым ядром.

  2. Выведенные атрибуты: переменные, используемые операторами внутри модуля. Имена атрибутов будут иметь префикс «inferred::». Это для того, чтобы отличить их от предопределённых атрибутов, полученных из аннотаций модуля Python. Выведенные атрибуты используются внутри подграфа локальной функции ONNX.

  • False (по умолчанию): экспорт вызовов nn.Module forward как узлов с высокой детализацией.
  • True: экспорт всех вызовов nn.Module forward как узлов локальной функции.

  • Set of type of nn.Module: export nn.Module forward calls as local function nodes,

    только если тип nn.Module присутствует в наборе.

Возможные исключения:

• torch.onnx.errors.CheckerError – Если проверка ONNX обнаруживает некорректную диаграмму ONNX.
• torch.onnx.errors.UnsupportedOperatorError – Если диаграмма ONNX не может быть экспортирована, потому что она использует оператор, который не поддерживается экспортером.
• torch.onnx.errors.OnnxExporterError – Другие ошибки, которые могут возникнуть во время экспорта. Все ошибки являются подклассами errors.OnnxExporterError.

torch.onnx.export_to_pretty_string(model, args, export_params=True, verbose=False, training=<TrainingMode.EVAL: 0>, input_names=None, output_names=None, operator_export_type=<OperatorExportTypes.ONNX: 0>, export_type=None, google_printer=False, opset_version=None, keep_initializers_as_inputs=None, custom_opsets=None, add_node_names=True, do_constant_folding=True, dynamic_axes=None) [source]

Аналогично export(), но возвращает текстовое представление модели ONNX. Различия в аргументах перечислены ниже. Все остальные аргументы такие же, как в export().

Параметры:

• add_node_names (bool, по умолчанию True) – Указывает, нужно ли устанавливать NodeProto.name. Это не влияет, если google_printer=True.
• google_printer (bool, по умолчанию False) – Если False, возвращает пользовательское, компактное представление модели. Если True, возвращает Message::DebugString() протокола protobuf, которое более подробное.

Возвращает:

UTF-8 строка, содержащая удобочитаемое представление модели ONNX.

torch.onnx.register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version) [source]

Регистрирует символическую функцию для пользовательского оператора.

Когда пользователь регистрирует символическую функцию для пользовательских/contrib операторов, рекомендуется добавить вывод формы для данного оператора с помощью API setType, иначе экспортированная диаграмма может иметь неправильный вывод формы в некоторых крайних случаях. Пример использования setType: test_aten_embedding_2 в test_operators.py.

См. раздел «Пользовательские операторы» в документации модуля для примера использования.

Параметры:

• symbolic_name (str) – Название пользовательского оператора в формате «<домен>::<оператор>».
• symbolic_fn (Callable) – Функция, принимающая диаграмму ONNX и аргументы входных данных текущего оператора и возвращающая новые узлы операторов для добавления в диаграмму.
• opset_version (int) – Версия ONNX opset, для которой регистрируется функция.

torch.onnx.unregister_custom_op_symbolic(symbolic_name, opset_version) [source]

Отменяет регистрацию symbolic_name.

См. раздел «Пользовательские операторы» в документации модуля для примера использования.

Параметры:

• symbolic_name (str) – Название пользовательского оператора в формате «<домен>::<оператор>».
• opset_version (int) – Версия ONNX opset, для которой отменяется регистрация.

torch.onnx.select_model_mode_for_export(model, mode) [source]

Менеджер контекста для временного изменения режима обучения model на mode, восстанавливая его при выходе из блока with.

Параметры:

• model – Тип и назначение такие же, как и аргумент model в export().
• mode (TrainingMode) – Тип и назначение такие же, как и аргумент training в export().

torch.onnx.is_in_onnx_export() [source]

Возвращает, находится ли процесс в середине экспорта ONNX.

Тип возвращаемого значения:

bool

torch.onnx.enable_log() [source]

Включает логирование ONNX.

torch.onnx.disable_log() [source]

Отключает логирование ONNX.

Классы