Spec-Zone.ru › PyTorch 2

Основанный на TorchScript экспортер ONNX

Примечание

Чтобы экспортировать модель ONNX с помощью TorchDynamo вместо TorchScript, см. torch.onnx.dynamo_export().

  • Пример: экспорт AlexNet из PyTorch в ONNX
  • Трейсинг против скриптинга

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

    • Избегайте NumPy и встроенных типов Python
    • Избегайте Tensor.data
    • Избегайте операций на месте при использовании tensor.shape в режиме трейсинга

  • Ограничения

    • Типы
    • Различия в реализации операторов

    • Неподдерживаемые шаблоны индексирования тензоров

      • Чтение / Получение
      • Запись / Установка

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

    • Внутренности экспортера ONNX

    • Операторы ATen

      • Список поддерживаемых операторов
      • Добавление поддержки оператора aten или квантованного оператора

    • torch.autograd.Functions

      • Статический символический метод
      • Встроенная функция автоградирования

    • Пользовательские операторы

      • ONNX-скрипт функции
      • Операторы C++
    • Обнаружение всех неконвертируемых операторов ATen сразу
  • Часто задаваемые вопросы

  • Python API

    • Функции
    • Классы

Пример: экспорт 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
import numpy as np

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 и записывает все операции, происходящие во время этого выполнения. Это означает, что если ваша модель динамична, например, меняет поведение в зависимости от входных данных, экспортированная модель не будет отображать это динамическое поведение. Мы рекомендуем проверить экспортированную модель и убедиться, что операторы выглядят разумно. Трейсинг будет развертывать циклы и условные операторы, экспортируя статическую граф, точно такую же, как и прослеженное выполнение. Если вы хотите экспортировать модель с динамическим управлением потоком, вам необходимо использовать скриптинг.

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

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

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

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

Избегайте NumPy и встроенных типов Python

Модели PyTorch могут быть написаны с использованием NumPy или типов и функций Python, но во время трейсинга любые переменные типов NumPy или Python (а не torch.Tensor) преобразуются в константы, что даст неправильный результат, если эти значения должны изменяться в зависимости от входных данных.

Например, вместо использования функций NumPy с numpy.ndarrays:

# 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)

Используйте операторы torch с torch.Tensors:

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

И вместо использования torch.Tensor.item() (которое преобразует Tensor в встроенное число Python):

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

Используйте поддержку torch для неявного преобразования одноэлементных тензоров:

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

Избегайте Tensor.data

Использование поля Tensor.data может привести к неверному трейсингу, а следовательно, и к неверной графе ONNX. Вместо этого используйте torch.Tensor.detach(). (Работа по полному удалению Tensor.data продолжается).

Избегайте операций на месте при использовании tensor.shape в режиме трейсинга

В режиме трейсинга формы, полученные из tensor.shape , отслеживаются как тензоры и делят одну и ту же память. Это может привести к несоответствию конечных значений вывода. В качестве обходного решения следует избегать использования операций на месте в этих сценариях. Например, в модели:

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 и seq_length делят одну и ту же память в режиме трейсинга. Этому можно было бы избежать, переписав операцию на месте:

real_seq_length = real_seq_length + 2

Ограничения

Типы

  • Поддерживаются только torch.Tensors, числовые типы, которые могут быть тривиально преобразованы в torch.Tensors (например, float, int), а также кортежи и списки этих типов в качестве входных или выходных данных модели. Входы и выходы Dict и str принимаются в режиме трейсинга, но:

    • Любая вычисление, зависящая от значения dict или str входных данных, будет заменена постоянным значением, увиденным во время единственного прослеженного выполнения.
    • Любой выходной dict будет молча заменён на уплощенный последовательность своих значений (ключи будут удалены). Например, {"foo": 1, "bar": 2} станет (1, 2).
    • Любой выходной str будет молча удален.
  • Некоторые операции с кортежами и списками не поддерживаются в режиме скриптинга из-за ограниченной поддержки ONNX для вложенных последовательностей. В частности, добавление кортежа к списку не поддерживается. В режиме трейсинга вложенные последовательности будут автоматически сплющены во время трейсинга.

Различия в реализации операторов

Из-за различий в реализации операторов запуск экспортированной модели в разных средах может давать разные результаты, как друг от друга, так и от PyTorch. Обычно эти различия незначительны, поэтому это должно волновать только в том случае, если ваше приложение чувствительно к этим небольшим различиям.

Неподдерживаемые шаблоны индексирования тензоров

Ниже перечислены шаблоны индексирования тензоров, которые не могут быть экспортированы. Если у вас возникают проблемы с экспортом модели, которая не содержит ни одного из нижеперечисленных неподдерживаемых шаблонов, пожалуйста, проверьте, что вы экспортируете с последней opset_version.

Чтение/Получение

При индексировании тензора для чтения не поддерживаются следующие шаблоны:

# 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!), вот как вы можете начать:

Внутренности экспортера ONNX

«Символическая функция» — это функция, которая декомпозирует оператор PyTorch в композицию нескольких операторов ONNX.

Во время экспорта каждый узел (содержащий оператор PyTorch) в графе TorchScript посещается экспортером в топологическом порядке. При посещении узла экспортер ищет зарегистрированные символические функции для этого оператора. Символические функции реализованы на Python. Символическая функция для операции с именем foo будет выглядеть примерно так:

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.
  """
  ...

Типы torch._C — это Python-обёртки вокруг типов, определённых на C++ в ir.h.

Процесс добавления символической функции зависит от типа оператора.

Операторы ATen

ATen — это встроенная библиотека тензоров PyTorch. Если оператор является оператором ATen (отображается в графе TorchScript с префиксом aten::), убедитесь, что он ещё не поддерживается.

Список поддерживаемых операторов

Посетите автоматически сгенерированный список поддерживаемых операторов TorchScript для получения подробной информации о том, какие операторы поддерживаются в каждом opset_version.

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

Если оператор не указан в списке выше:

  • Определите символическую функцию в torch/onnx/symbolic_opset<version>.py, например, torch/onnx/symbolic_opset9.py. Убедитесь, что функция имеет то же имя, что и функция ATen, которая может быть объявлена в torch/_C/_VariableFunctions.pyi или torch/nn/functional.pyi (эти файлы генерируются на этапе сборки, поэтому не появятся в вашей копировании, пока вы не соберете PyTorch).
  • По умолчанию первый аргумент — граф ONNX. Другие имена аргументов должны ТОЧНО совпадать с именами в файле .pyi, так как распределение выполняется с использованием именованных аргументов.
  • В символической функции, если оператор находится в стандартном наборе операторов ONNX, нам нужно только создать узел для представления оператора ONNX в графе. В противном случае, мы можем составить несколько стандартных операторов, которые имеют эквивалентную семантику оператора ATen.

Вот пример обработки отсутствия символической функции для оператора ELU.

Если мы выполним следующий код:

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)

Поскольку мы видим aten::elu в графе, мы знаем, что это оператор ATen.

Мы проверяем список операторов ONNX и подтверждаем, что Elu стандартизирован в ONNX.

Мы находим сигнатуру для elu в torch/nn/functional.pyi. 

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

Мы добавляем следующие строки в symbolic_opset9.py:

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

Теперь PyTorch может экспортировать модели, содержащие оператор aten::elu!

См. файлы torch/onnx/symbolic_opset*.py для получения дополнительных примеров.

torch.autograd.Functions

Если оператор является подклассом torch.autograd.Function, существуют три способа его экспорта.

Статический символический метод

Вы можете добавить статический метод с именем symbolic в ваш класс функции. Он должен возвращать операторы ONNX, представляющие поведение функции в ONNX. Например:

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)))

Встроенная функция Autograd

В тех случаях, когда статический символический метод не предоставляется для последующего torch.autograd.Function или когда функция для регистрации prim::PythonOp в качестве пользовательских символических функций не предоставлена, torch.onnx.export() пытается встроить граф, соответствующий этому torch.autograd.Function, таким образом, эта функция разбивается на отдельные операторы, которые использовались внутри функции. Экспорт должен быть успешным, если эти отдельные операторы поддерживаются. Например:

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)

Если вам нужно избежать встраивания torch.autograd.Function, вы должны экспортировать модели с operator_export_type установленными на ONNX_FALLTHROUGH или ONNX_ATEN_FALLBACK.

Пользовательские операторы

Вы можете экспортировать свою модель с пользовательскими операторами, которые включают комбинацию многих стандартных операторов ONNX или управляются собственным определённым C++ бэкендом.

Функции ONNX-script

Если оператор не является стандартным оператором ONNX, но может быть составлен из нескольких существующих операторов ONNX, вы можете использовать ONNX-script для создания внешней функции ONNX для поддержки оператора. Вы можете экспортировать его, следуя этому примеру:

import onnxscript
# There are three opset version needed to be aligned
# This is (1) the opset version in ONNX function
from onnxscript.onnx_opset import opset15 as op
opset_version = 15

x = torch.randn(1, 2, 3, 4, requires_grad=True)
model = torch.nn.SELU()

custom_opset = onnxscript.values.Opset(domain="onnx-script", version=1)

@onnxscript.script(custom_opset)
def Selu(X):
    alpha = 1.67326  # auto wrapped as Constants
    gamma = 1.0507
    alphaX = op.CastLike(alpha, X)
    gammaX = op.CastLike(gamma, X)
    neg = gammaX * (alphaX * op.Exp(X) - alphaX)
    pos = gammaX * X
    zero = op.CastLike(0, X)
    return op.Where(X <= zero, neg, pos)

# setType API provides shape/type to ONNX shape/type inference
def custom_selu(g: jit_utils.GraphContext, X):
    return g.onnxscript_op(Selu, X).setType(X.type())

# Register custom symbolic function
# There are three opset version needed to be aligned
# This is (2) the opset version in registry
torch.onnx.register_custom_op_symbolic(
    symbolic_name="aten::selu",
    symbolic_fn=custom_selu,
    opset_version=opset_version,
)

# There are three opset version needed to be aligned
# This is (2) the opset version in exporter
torch.onnx.export(
    model,
    x,
    "model.onnx",
    opset_version=opset_version,
    # only needed if you want to specify an opset version > 1.
    custom_opsets={"onnx-script": 2}
)

В приведенном выше примере он экспортируется как пользовательский оператор в наборе операторов «onnx-script». При экспорте пользовательского оператора вы можете указать версию пользовательского домена, используя словарь custom_opsets при экспорте. Если не указано иное, версия набора операторов по умолчанию составляет 1.

ПРИМЕЧАНИЕ: Обратите внимание на соответствие версии набора операторов, упомянутой в приведенном выше примере, и убедитесь, что они используются на шаге экспорта. Пример использования функции onnx-script — это бета-версия с точки зрения активного развития onnx-script. Пожалуйста, следуйте последней версии ONNX-script

Операторы C++

Если модель использует пользовательский оператор, реализованный на C++ так, как описано в Расширение TorchScript с пользовательскими операторами C++, вы можете экспортировать его, следуя этому примеру:

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}
)

В примере выше он экспортируется как пользовательский оператор в наборе операторов «custom_domain». При экспорте пользовательского оператора вы можете указать версию пользовательского домена, используя словарь custom_opsets при экспорте. Если не указано иное, версия набора операторов по умолчанию составляет 1.

Интерпретатор, использующий модель, должен поддерживать пользовательский оператор. См. Пользовательские операторы Caffe2, Пользовательские операторы ONNX Runtime или документацию вашего интерпретатора по выбору.

Обнаружение всех необратимых операторов ATen сразу

Когда экспорт завершается ошибкой из-за необратимого оператора ATen, может быть больше одного такого оператора, но сообщение об ошибке упоминает только первый. Чтобы обнаружить все необратимые операторы сразу, вы можете:

# 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))

Множество приблизительное, так как некоторые операторы могут быть удалены во время процесса преобразования и не требуют преобразования. Некоторые другие операторы могут иметь частичную поддержку, которая приведёт к ошибке преобразования с определёнными входными данными, но это должно дать вам общее представление о том, какие операторы не поддерживаются. Не стесняйтесь открывать GitHub-вопросы по запросам поддержки операторов.

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

В: Я экспортировал свою модель 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 в opset 11.

Python API

Функции

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, autograd_inlining=True) [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) является допустимым вызовом модели. Все аргументы, не являющиеся тензорами, будут жестко закодированы в экспортированной модели; все тензорные аргументы станут входными данными экспортированной модели в порядке их появления в кортеже.

    1. ТЕНЗОР:

      args = torch.Tensor([1])

    Эквивалентно кортежу из одного тензора.

    1. КОРТЕЖ АРГУМЕНТОВ С ЗАКЛЮЧИТЕЛЬНЫМ СЛОВАРЕМ ИМЕНОВАННЫХ АРГУМЕНТОВ:

      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

      (в домене стандартных операций).

    • OperatorExportTypes.ONNX_FALLTHROUGH: Try to convert all ops

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

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

      экспортируются как операции ATen (в домене операций «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 (целое число, по умолчанию 17) – версия доменных операций (ai.onnx) для мишени. Должно быть >= 7 и <= 17.
  • 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, то инициализаторы не добавляются как входные данные в графу, и в качестве входных данных добавляются только непараметрические входные данные. Это может позволить лучшую оптимизацию (например, константное складывание) бэкендами/средами выполнения.

    Если True, этап deduplicate_initializers не будет выполнен. Это означает, что инициализаторы с дублированными значениями не будут дедуплицированы и будут обрабатываться как отдельные входные данные в графу. Это позволяет передавать различные входные инициализаторы во время выполнения после экспорта.

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

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

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

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

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

  • custom_opsets (dict[str, int], по умолчанию пустой словарь) –

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

    • КЛЮЧ (str): имя домена опсета
    • ЗНАЧЕНИЕ (int): версия опсета

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

  • export_modules_as_functions (bool или множество из тип класса 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 найден в множестве.

  • autograd_inlining (bool, по умолчанию True) – Флаг, используемый для управления включением/выключением инлайнинга функций autograd. Подробности см. в https://github.com/pytorch/pytorch/pull/74765.
Исключения
  • 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. Это не имеет значения, если 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.

torch.onnx.verification.find_mismatch(model, input_args, do_constant_folding=True, training=<TrainingMode.EVAL: 0>, opset_version=None, keep_initializers_as_inputs=True, verbose=False, options=None) [source]

Найти все несоответствия между исходной моделью и экспортированной моделью.

Экспериментальная функция. API может быть изменён.

Этот инструмент помогает отлаживать несоответствия между исходной моделью PyTorch и экспортированной моделью ONNX. Он выполняет двоичный поиск в графе модели, чтобы найти минимальную подграф, демонстрирующую несоответствие.

Параметры
  • model (Union[Module, ScriptModule]) – Модель, подлежащая экспорту.
  • input_args (Tuple[Any, ...]) – Входные аргументы модели.
  • do_constant_folding (bool) – Аналогично do_constant_folding в torch.onnx.export().
  • training (TrainingMode) – Аналогично training в torch.onnx.export().
  • opset_version (Optional[int]) – Аналогично opset_version в torch.onnx.export().
  • keep_initializers_as_inputs (bool) – Аналогично keep_initializers_as_inputs в torch.onnx.export().
  • verbose (bool) – Аналогично verbose в torch.onnx.export().
  • options (Optional[VerificationOptions]) – Параметры для проверки несоответствий.
Возвращает

Объект GraphInfo, содержащий информацию о несоответствии.

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

Пример:

>>> import torch
>>> import torch.onnx.verification
>>> torch.manual_seed(0)
>>> opset_version = 15
>>> # Define a custom symbolic function for aten::relu.
>>> # The custom symbolic function is incorrect, which will result in mismatches.
>>> def incorrect_relu_symbolic_function(g, self):
...     return self
>>> torch.onnx.register_custom_op_symbolic(
...     "aten::relu",
...     incorrect_relu_symbolic_function,
...     opset_version=opset_version,
... )
>>> class Model(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.layers = torch.nn.Sequential(
...             torch.nn.Linear(3, 4),
...             torch.nn.ReLU(),
...             torch.nn.Linear(4, 5),
...             torch.nn.ReLU(),
...             torch.nn.Linear(5, 6),
...         )
...     def forward(self, x):
...         return self.layers(x)
>>> graph_info = torch.onnx.verification.find_mismatch(
...     Model(),
...     (torch.randn(2, 3),),
...     opset_version=opset_version,
... )
===================== Mismatch info for graph partition : ======================
================================ Mismatch error ================================
Tensor-likes are not close!
Mismatched elements: 12 / 12 (100.0%)
Greatest absolute difference: 0.2328854203224182 at index (1, 2) (up to 1e-07 allowed)
Greatest relative difference: 0.699536174352349 at index (1, 3) (up to 0.001 allowed)
==================================== Tree: =====================================
5 X   __2 X    __1 ✓
id:  |  id: 0 |  id: 00
     |        |
     |        |__1 X (aten::relu)
     |           id: 01
     |
     |__3 X    __1 ✓
        id: 1 |  id: 10
              |
              |__2 X     __1 X (aten::relu)
                 id: 11 |  id: 110
                        |
                        |__1 ✓
                           id: 111
=========================== Mismatch leaf subgraphs: ===========================
['01', '110']
============================= Mismatch node kinds: =============================
{'aten::relu': 2}

Классы

JitScalarType

Типы скаляров, определённые в torch.
torch.onnx.verification.GraphInfo

GraphInfo содержит информацию о валидации графа TorchScript и его преобразованного графа ONNX.
torch.onnx.verification.VerificationOptions

Параметры для проверки экспорта ONNX.

© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/onnx_torchscript.html

Spec-Zone.ru

Настройки Оффлайн Что нового Помощь О нас
Spec-Zone .ru
спецификации, руководства, описания, API