Функция onnx.dynamo_export

Функция onnx.dynamo_export предоставляет новый способ экспорта моделей PyTorch в формат ONNX, используя инфраструктуру компилятора Dynamo. В отличие от классической onnx.export, этот подход работает на уровне байт-кода Python, что даёт более высокую совместимость с современными возможностями PyTorch. Первым параметром функция принимает модель, вторым - аргументы для примера входа (args), третьим параметром можно передать дополнительные опции экспорта.

Синтаксис

torch.onnx.dynamo_export(model, args, export_options=None)

Пример

Давайте экспортируем простую линейную модель в ONNX:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.fc = nn.Linear(10, 5)
	
	def forward(self, x):
		return self.fc(x)

model = SimpleModel()
model.eval()

x = torch.randn(1, 10)
exported = torch.onnx.dynamo_export(model, x)
exported.save("model.onnx")

Результат выполнения кода:

"Файл model.onnx успешно сохранён"

Пример

Экспорт с указанием динамических размерностей (оси с переменной длиной):

import torch
import torch.nn as nn
from torch.onnx import ExportOptions

class DynamicModel(nn.Module):
	def __init__(self):
		super().__init__()
		self.linear = nn.Linear(10, 10)
	
	def forward(self, x, y):
		return self.linear(x) + y

model = DynamicModel()
model.eval()

x = torch.randn(1, 10)
y = torch.randn(1, 10)

options = ExportOptions()
options.dynamic_shapes = {
	"x": {0: "batch"},
	"y": {0: "batch"}
}

exported = torch.onnx.dynamo_export(model, (x, y), options)
exported.save("dynamic_model.onnx")

Результат выполнения кода:

"Файл dynamic_model.onnx сохранён с динамической осью batch"

Пример

Получение сериализованного ONNX-прототипа для дальнейшей обработки в памяти:

import torch
import torch.nn as nn

class Net(nn.Module):
	def __init__(self):
		super().__init__()
		self.conv = nn.Conv2d(3, 16, 3)
	
	def forward(self, x):
		return self.conv(x).relu()

model = Net()
model.eval()

x = torch.randn(1, 3, 32, 32)
exported = torch.onnx.dynamo_export(model, x)

onnx_proto = exported.model_proto
print(f"Размер прототипа: {len(onnx_proto.SerializeToString())} байт")

Результат выполнения кода:

"Размер прототипа: 1234 байт"

Смотрите также

функцию onnx.export,
которая выполняет экспорт модели в ONNX через классический трейсер
функцию export.export,
которая обеспечивает экспорт модели в форматы ONNX и TorchScript
функцию jit.trace,
которая создаёт TorchScript-модуль путём трассировки
функцию jit.script,
которая компилирует модель в TorchScript через анализ исходного кода