Справочник API torch.export

Создано: 17 июля 2025 г. | Последнее обновление: 24 апреля 2026 г.

torch.export.export(mod, args, kwargs=None, *, dynamic_shapes=None, strict=False, preserve_module_call_signature=(), prefer_deferred_runtime_asserts_over_guards=False) [исходный код]

export() принимает любой nn.Module вместе с примерами входных данных и создает трассированный граф, представляющий только вычисления Tensor этой функции, в режиме Ahead-of-Time (AOT); впоследствии его можно выполнять с другими входными данными или сериализовать. Трассированный граф (1) создает нормализованные операторы из функционального набора операторов ATen (а также любые пользовательские операторы), (2) устраняет весь поток управления Python и структуры данных (за некоторыми исключениями) и (3) фиксирует набор ограничений на формы, необходимых для подтверждения корректности этой нормализации и устранения потока управления для будущих входных данных.

Гарантия корректности

Во время трассировки export() фиксирует предположения о формах, сделанные пользовательской программой и базовыми ядрами операторов PyTorch. Выходной объект ExportedProgram считается допустимым только при выполнении этих предположений.

При трассировке делаются предположения о формах (а не значениях) входных тензоров. Для успешного выполнения export() эти предположения необходимо проверить во время захвата графа. В частности:

Предположения о статических формах входных тензоров проверяются автоматически без дополнительных действий.
Предположения о динамических формах входных тензоров требуют явного указания с помощью API Dim() для создания динамических измерений и их связывания с примерами входных данных через аргумент dynamic_shapes.

Если какое-либо предположение не удается проверить, возникает критическая ошибка. В этом случае сообщение об ошибке будет содержать предлагаемые изменения спецификации, необходимые для проверки предположений. Например, export() может предложить следующее изменение определения динамического измерения dim0_x, например встречающегося в форме, связанной с входом x, которое ранее было определено как Dim("dim0_x"):

dim = Dim("dim0_x", max=5)

Этот пример означает, что сгенерированному коду для корректной работы требуется, чтобы измерение 0 входа x было меньше или равно 5. Вы можете изучить предлагаемые изменения определений динамических измерений, а затем дословно скопировать их в свой код, не меняя аргумент dynamic_shapes в вызове export().

Параметры:

mod (Module) – Будет трассирован метод forward этого модуля.
args (tuple[Any, ...]) – Примеры позиционных входных данных.
kwargs (Mapping[str, Any] | None) – Необязательные примеры входных данных, передаваемых по именам.
dynamic_shapes (Any | None) –

Необязательный аргумент, типом которого может быть: 1) словарь, сопоставляющий имена аргументов f их спецификациям динамических форм. Для аргумента **kwargs словарь может использовать фактические имена ключевых слов, переданные в f, если имя ключевого слова не совпадает с именем другого входного аргумента. 2) кортеж, задающий спецификации динамических форм для каждого входного аргумента в исходном порядке. Если вы указываете динамичность для аргументов-ключевых слов, передавать их нужно в порядке, заданном исходной сигнатурой функции.

Динамическую форму аргумента-тензора можно указать как (1) словарь, сопоставляющий индексы динамических измерений типам Dim(); включать в этот словарь индексы статических измерений не требуется, но если они включены, им нужно сопоставить None; либо как (2) кортеж / список типов Dim() или None, где типы Dim() соответствуют динамическим измерениям, а статические измерения обозначаются значением None. Аргументы, представляющие собой словари или кортежи / списки тензоров, задаются рекурсивно с помощью отображений или последовательностей содержащихся в них спецификаций.

API ShapesSpec. dynamic_shapes также может быть объектом torch.fx.experimental.dynamic_spec.ShapesSpec (или его сокращенной формой torch.fx.experimental.dynamic_spec.ParamsSpec). Это новый унифицированный API без привязки к конкретной реализации для compile, pre-compile, export и т. д.; в дальнейшем рекомендуется использовать его для задания динамических форм при экспорте. Это тот же API спецификаций, который предоставляется через dynamic_shapes= в torch.compile().

Ключи ParamsSpec — это имена параметров вызываемого объекта, который трассируется (для nn.Module — параметры forward); аргументы-ключевые слова и **kwargs также указываются по имени.

Основные свойства (подробности см. в torch.fx.experimental.dynamic_spec):
- Только без привязки. Измерения / скаляры, отмеченные как динамические, становятся SymInt без привязки (символами u) и никогда не специализируются (в том числе без специализации 0/1).
- Предположения и производные выражения. Измерение может быть выражением над символами спецификации (например, TensorSpec([B * 2, ...])), а также можно передавать реляционные assumptions между символами (например, [B % 2 == 0]), которые проверяет export.
- Без незаметной специализации. Гарантируется, что экспортированный граф будет корректен при выполнении всех заданных предположений; в противном случае экспорт завершится ошибкой. (В отличие от него, Dim.DYNAMIC / Dim.AUTO могут незаметно специализировать динамическое измерение, создавая граф, корректный не для всех входных данных.)
Пример:
```
batch = ShapeVar("batch", min=2, max=128)
ep = torch.export.export(
    mod,
    (torch.randn(8, 3), torch.randn(16, 3)),
    dynamic_shapes=ShapesSpec(
        params=ParamsSpec(
            {
                "x": TensorSpec([batch, 3]),
                "y": TensorSpec([batch * 2, 3]),  # derived expression
            }
        ),
        assumptions=[batch % 2 == 0],
    ),
    strict=True,
)
```
strict (bool) – Если параметр отключен (по умолчанию), функция export трассирует программу во время выполнения Python, что само по себе не проверяет некоторые неявные предположения, заложенные в граф. При этом проверяются наиболее важные предположения, например безопасность форм. Если параметр включен (с помощью strict=True), функция export трассирует программу через TorchDynamo, который обеспечивает корректность результирующего графа. Поддержка возможностей Python в TorchDynamo ограничена, поэтому может возникнуть больше ошибок. Обратите внимание: переключение этого аргумента не меняет спецификацию результирующего IR, и модель будет сериализована одинаково независимо от переданного значения.
preserve_module_call_signature (tuple[str, ...]) – Список путей подмодулей, для которых исходные соглашения о вызовах сохраняются в виде метаданных. Эти метаданные используются при вызове torch.export.unflatten для сохранения исходных соглашений о вызовах модулей.

Возвращает:

Объект ExportedProgram, содержащий трассированный вызываемый объект.

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

ExportedProgram

Допустимые типы входных и выходных данных

Допустимые типы входных данных (для args и kwargs) и выходных данных:

Примитивные типы, то есть torch.Tensor, int, float, bool и str.
Классы данных, но сначала их необходимо зарегистрировать, вызвав register_dataclass().
Вложенные структуры данных, состоящие из dict, list, tuple, namedtuple и OrderedDict, содержащих все перечисленные выше типы.

class torch.export.ExportedProgram(root, graph, graph_signature, state_dict, range_constraints, module_call_graph, example_inputs=None, constants=None, *, verifiers=None) [исходный код]

Пакет программы, созданной с помощью export(). Он содержит torch.fx.Graph, представляющий вычисления Tensor, state_dict со значениями тензоров всех поднятых параметров и буферов, а также различные метаданные.

Можно вызывать ExportedProgram так же, как исходный вызываемый объект, трассированный с помощью export(), соблюдая те же соглашения о вызовах.

Для преобразования графа используйте свойство .module, чтобы получить доступ к torch.fx.GraphModule. Затем можно использовать преобразование FX для переписывания графа. После этого достаточно снова вызвать export(), чтобы создать корректный ExportedProgram.

buffers() [исходный код]

Возвращает итератор по буферам исходного модуля.

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

Этот API является экспериментальным и НЕ обратно совместим.

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

property call_spec: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property constants: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property dialect: str: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property example_inputs: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property graph: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property graph_module: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property graph_signature: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

module(check_guards=True) [исходный код]

Возвращает автономный GraphModule со встроенными параметрами/буферами.

Если задано check_guards=True (по умолчанию), создается подмодуль _guards_fn, а сразу после заполнителей в граф вставляется вызов подмодуля _guards_fn. Этот модуль проверяет входные данные на соответствие защитным условиям.
Если задано check_guards=False, часть этих проверок выполняется предварительным хуком forward для модуля графа. Подмодуль _guards_fn не создается.

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

property module_call_graph: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

named_buffers() [исходный код]

Возвращает итератор по буферам исходного модуля, выдавая как имя буфера, так и сам буфер.

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

Этот API является экспериментальным и НЕ обратно совместим.

Тип возвращаемого значения:: Iterator[tuple[str, Tensor]]

named_parameters() [исходный код]

Возвращает итератор по параметрам исходного модуля, выдавая как имя параметра, так и сам параметр.

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

Этот API является экспериментальным и НЕ обратно совместим.

Тип возвращаемого значения:: Iterator[tuple[str, Parameter]]

parameters() [исходный код]

Возвращает итератор по параметрам исходного модуля.

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

Этот API является экспериментальным и НЕ обратно совместим.

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

property range_constraints: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

run_decompositions(decomp_table=None, decompose_custom_triton_ops=False) [исходный код]

Выполняет набор декомпозиций экспортированной программы и возвращает новую экспортированную программу. По умолчанию выполняются декомпозиции Core ATen, чтобы получить операторы из набора операторов Core ATen.

Пока что совместные графы не декомпозируются.

Параметры:: decomp_table (dict[OperatorBase, Callable] | None) – Необязательный аргумент, задающий поведение декомпозиции для операторов Aten: (1) если значение None, выполняется декомпозиция до базовых декомпозиций aten; (2) если значение пустое, декомпозиция операторов не выполняется.
Тип возвращаемого значения:: ExportedProgram

Примеры:

Если декомпозиция не требуется

ep = torch.export.export(model, ...)
ep = ep.run_decompositions(decomp_table={})

Если нужно получить набор операторов Core ATen, кроме некоторых операторов, можно сделать следующее:

ep = torch.export.export(model, ...)
decomp_table = torch.export.default_decompositions()
decomp_table[your_op] = your_custom_decomp
ep = ep.run_decompositions(decomp_table=decomp_table)

property state_dict: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property tensor_constants: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

validate() [исходный код]: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property verifier: Any: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

property verifiers: Предупреждение

Этот API является экспериментальным и НЕ обратно совместим.

class torch.export.dynamic_shapes.AdditionalInputs [исходный код]

Выводит dynamic_shapes на основе дополнительных входных данных.

Примечание: этот метод формирует только формат Dim на основе dynamic_shapes; он пока не создает структурированный API ShapesSpec / ParamsSpec.

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

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

Пример:

args0, kwargs0 = ...  # example inputs for export

# other representative inputs that the exported program will run on
dynamic_shapes = torch.export.AdditionalInputs()
dynamic_shapes.add(args1, kwargs1)
...
dynamic_shapes.add(argsN, kwargsN)

torch.export(..., args0, kwargs0, dynamic_shapes=dynamic_shapes)

add(args, kwargs=None) [исходный код]: Добавляет вход args() и kwargs().

dynamic_shapes(m, args, kwargs=None) [исходный код]: Выводит структуру pytree dynamic_shapes(), объединяя формы исходных входных данных args() и kwargs() с формами каждого дополнительного аргумента args и kwargs.

verify(ep) [исходный код]: Проверяет допустимость экспортированной программы для каждого дополнительного входного аргумента.

class torch.export.dynamic_shapes.Dim(name, *, min=None, max=None) [исходный код]

Класс Dim позволяет пользователям задавать динамичность в экспортируемых программах. Помечая измерение с помощью Dim, компилятор связывает его с символическим целым числом, имеющим динамический диапазон.

API можно использовать двумя способами: подсказки Dim (то есть автоматические динамические формы: Dim.AUTO, Dim.DYNAMIC, Dim.STATIC) или именованные Dim (то есть Dim("name", min=1, max=2)).

Подсказки Dim упрощают экспорт: пользователю нужно лишь указать, является ли измерение динамическим, статическим или решение остается за компилятором (Dim.AUTO). Во время экспорта компилятор автоматически выводит остальные ограничения на минимальный/максимальный диапазон и взаимосвязи между измерениями.

Пример:

class Foo(nn.Module):
    def forward(self, x, y):
        assert x.shape[0] == 4
        assert y.shape[0] >= 16
        return x @ y


x = torch.randn(4, 8)
y = torch.randn(8, 16)
dynamic_shapes = {
    "x": {0: Dim.AUTO, 1: Dim.AUTO},
    "y": {0: Dim.AUTO, 1: Dim.AUTO},
}
ep = torch.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)

В этом случае export вызовет исключение, если заменить все вхождения Dim.AUTO на Dim.DYNAMIC, поскольку модель накладывает на x.shape[0] ограничение статичности.

Компилятор также может сгенерировать для более сложных взаимосвязей между измерениями узлы утверждений времени выполнения, например (x.shape[0] + y.shape[1]) % 4 == 0; они будут вызваны, если входные данные во время выполнения не удовлетворяют таким ограничениям.

Для подсказок Dim также можно задать минимальные и максимальные границы, например Dim.AUTO(min=16, max=32), Dim.DYNAMIC(max=64); компилятор выведет остальные ограничения в пределах диапазонов. Исключение будет вызвано, если допустимый диапазон целиком выходит за указанный пользователем диапазон.

Именованные Dim позволяют задавать динамичность более строго: исключения вызываются, если выведенные компилятором ограничения не соответствуют спецификации пользователя. Например, при экспорте предыдущей модели пользователю нужно было бы указать следующий аргумент dynamic_shapes:

s0 = Dim("s0")
s1 = Dim("s1", min=16)
dynamic_shapes = {
    "x": {0: 4, 1: s0},
    "y": {0: s0, 1: s1},
}
ep = torch.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)

Именованные Dim также позволяют задавать взаимосвязи между измерениями вплоть до одномерных линейных зависимостей. Например, следующее выражение указывает, что одно измерение кратно другому с прибавлением 4:

s0 = Dim("s0")
s1 = 3 * s0 + 4

class torch.export.dynamic_shapes.ShapesCollection [исходный код]

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

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

Пример:

args = {"x": tensor_x, "others": [tensor_y, tensor_z]}

dim = torch.export.Dim(...)
dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[tensor_y] = {0: dim * 2}
# This is equivalent to the following (now auto-generated):
# dynamic_shapes = {"x": (dim, dim + 1, 8), "others": [{0: dim * 2}, None]}

torch.export(..., args, dynamic_shapes=dynamic_shapes)

Чтобы задать динамичность для целых чисел, сначала нужно обернуть их с помощью _IntWrapper, чтобы получить для каждого целого числа «уникальный идентификатор».

Пример:

args = {"x": tensor_x, "others": [int_x, int_y]}
# Wrap all ints with _IntWrapper
mapped_args = pytree.tree_map_only(int, lambda a: _IntWrapper(a), args)

dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[mapped_args["others"][0]] = Dim.DYNAMIC

# This is equivalent to the following (now auto-generated):
# dynamic_shapes = {"x": (dim, dim + 1, 8), "others": [Dim.DYNAMIC, None]}

torch.export(..., args, dynamic_shapes=dynamic_shapes)

dynamic_shapes(m, args, kwargs=None) [исходный код]: Создает структуру pytree dynamic_shapes() в соответствии с args() и kwargs().

torch.export.dynamic_shapes.dims(*names, min=None, max=None) [исходный код]

Вспомогательная функция для создания нескольких типов Dim().

Возвращает:: Кортеж типов Dim().
Тип возвращаемого значения:: tuple[Dim, …]

torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes(msg, dynamic_shapes) [исходный код]

При экспорте с помощью dynamic_shapes() экспорт может завершиться ошибкой ConstraintViolation, если спецификация не соответствует ограничениям, выведенным при трассировке модели. В сообщении об ошибке могут предлагаться исправления — изменения, которые можно внести в dynamic_shapes() для успешного экспорта.

Пример сообщения об ошибке ConstraintViolation:

Suggested fixes:

    dim = Dim('dim', min=3, max=6)  # this just refines the dim's range
    dim = 4  # this specializes to a constant
    dy = dx + 1  # dy was specified as an independent dim, but is actually tied to dx with this relation

Эта вспомогательная функция принимает сообщение об ошибке ConstraintViolation и исходную спецификацию dynamic_shapes() и возвращает новую спецификацию dynamic_shapes() с учётом предложенных исправлений.

Пример использования:

try:
    ep = export(mod, args, dynamic_shapes=dynamic_shapes)
except torch._dynamo.exc.UserError as exc:
    new_shapes = refine_dynamic_shapes_from_suggested_fixes(
        exc.msg, dynamic_shapes
    )
    ep = export(mod, args, dynamic_shapes=new_shapes)

Тип возвращаемого значения:: dict[str, Any] | tuple[Any] | list[Any]

torch.export.save(ep, f, *, extra_files=None, opset_version=None, pickle_protocol=2) [исходный код]

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

Функциональность находится в активной разработке; сохранённые файлы могут быть несовместимы с более новыми версиями PyTorch.

Сохраняет ExportedProgram в файловый объект. Затем его можно загрузить с помощью API Python torch.export.load.

Параметры:

ep (ExportedProgram) – Экспортированная программа, которую нужно сохранить.
f (str | os.PathLike[str] | IO[bytes]) – объект, реализующий методы write и flush, или строка с именем файла.
extra_files (Optional[Dict[str, Any]]) – Сопоставление имён файлов с содержимым, которое будет сохранено вместе с f.
opset_version (Optional[Dict[str, int]]) – Сопоставление имён наборов операторов с версиями этих наборов.
pickle_protocol (int) – Позволяет переопределить протокол по умолчанию.

Пример:

import torch
import io


class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10


ep = torch.export.export(MyModule(), (torch.randn(5),))

# Save to file
torch.export.save(ep, "exported_program.pt2")

# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.export.save(ep, buffer)

# Save with extra files
extra_files = {"foo.txt": b"bar".decode("utf-8")}
torch.export.save(ep, "exported_program.pt2", extra_files=extra_files)

torch.export.load(f, *, extra_files=None, expected_opset_version=None) [исходный код]

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

torch.export.load() использует pickle для загрузки моделей. Никогда не загружайте данные из ненадёжного источника.

Загружает ExportedProgram, ранее сохранённый с помощью torch.export.save.

Параметры:

f (str | os.PathLike[str] | IO[bytes]) – файловый объект (должен реализовывать методы write и flush) или строка с именем файла.
extra_files (Optional[Dict[str, Any]]) – Дополнительные имена файлов из этого сопоставления будут загружены, а их содержимое будет сохранено в переданном сопоставлении.
expected_opset_version (Optional[Dict[str, int]]) – Сопоставление имён наборов операторов с ожидаемыми версиями этих наборов.

Возвращает:

Объект ExportedProgram

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

ExportedProgram

Пример:

import torch
import io

# Load ExportedProgram from file
ep = torch.export.load("exported_program.pt2")

# Load ExportedProgram from io.BytesIO object
with open("exported_program.pt2", "rb") as f:
    buffer = io.BytesIO(f.read())
buffer.seek(0)
ep = torch.export.load(buffer)

# Load with extra files.
extra_files = {"foo.txt": ""}  # values will be replaced with data
ep = torch.export.load("exported_program.pt2", extra_files=extra_files)
print(extra_files["foo.txt"])
print(ep(torch.randn(5)))

torch.export.pt2_archive._package.package_pt2(f, *, exported_programs=None, aoti_files=None, extra_files=None, opset_version=None, pickle_protocol=2, executorch_files=None) [исходный код]

Сохраняет артефакты в формате PT2Archive. Затем артефакт можно загрузить с помощью load_pt2.

Параметры:

f (str | os.PathLike[str] | IO[bytes]) – файловый объект (должен реализовывать методы write и flush) или строка с именем файла.
exported_programs (Union[ExportedProgram, dict[str, ExportedProgram]]) – Экспортированная программа, которую нужно сохранить, или словарь, сопоставляющий имена моделей с экспортированными программами. Экспортированная программа будет сохранена в models/*.json. Если указана только одна ExportedProgram, ей автоматически присваивается имя «model».
aoti_files (Union[list[str], dict[str, list[str]]]) – Список файлов, созданных AOTInductor с помощью torch._inductor.aot_compile(..., {"aot_inductor.package": True}), или словарь, сопоставляющий имена моделей с файлами, созданными AOTInductor. Если указан только один набор файлов, ему автоматически присваивается имя «model».
extra_files (Optional[Dict[str, Any]]) – Сопоставление имён файлов с содержимым, которое будет сохранено в составе pt2.
opset_version (Optional[Dict[str, int]]) – Сопоставление имён наборов операторов с версиями этих наборов.
pickle_protocol (int) – Позволяет переопределить протокол по умолчанию.
executorch_files (Optional[dict[str, bytes]]) – Необязательные артефакты executorch для сохранения.

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

str | PathLike[str] | IO[bytes]

torch.export.pt2_archive._package.load_pt2(f, *, expected_opset_version=None, run_single_threaded=False, num_runners=1, device_index=-1, load_weights_from_disk=False) [исходный код]

Загружает все артефакты, ранее сохранённые с помощью package_pt2.

Параметры:

f (str | os.PathLike[str] | IO[bytes]) – файловый объект (должен реализовывать методы write и flush) или строка с именем файла.
expected_opset_version (Optional[Dict[str, int]]) – Сопоставление имён наборов операторов с ожидаемыми версиями этих наборов.
num_runners (int) – Количество исполнителей для загрузки артефактов AOTInductor.
run_single_threaded (bool) – Нужно ли запускать модель без логики синхронизации потоков. Это полезно для предотвращения конфликтов с CUDAGraphs.
device_index (int) – Индекс устройства, на которое следует загрузить пакет PT2. По умолчанию используется device_index=-1, соответствующий устройству cuda при использовании CUDA. Например, передача device_index=1 загрузит пакет на cuda:1.

Возвращает:

Объект PT2ArchiveContents, содержащий все объекты из PT2.

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

PT2ArchiveContents

torch.export.draft_export(mod, args, kwargs=None, *, dynamic_shapes=None, preserve_module_call_signature=(), strict=False, prefer_deferred_runtime_asserts_over_guards=False) [исходный код]

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

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

class torch.export.unflatten.FlatArgsAdapter [исходный код]

Адаптирует входные аргументы с помощью input_spec для согласования с target_spec.

abstract adapt(target_spec, input_spec, input_args, metadata=None, obj=None) [исходный код]

ПРИМЕЧАНИЕ: Этот адаптер может изменять переданный input_args_with_path.

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

get_flat_arg_paths() [исходный код]

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

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

class torch.export.unflatten.InterpreterModule(graph, ty=None) [исходный код]: Модуль, который использует torch.fx.Interpreter для выполнения вместо обычной генерации кода, применяемой GraphModule. Это обеспечивает более подробную информацию в трассировке стека и упрощает отладку выполнения.

class torch.export.unflatten.InterpreterModuleDispatcher(attrs, call_modules) [исходный код]: Модуль, содержащий последовательность InterpreterModule, соответствующих последовательности вызовов этого модуля. Каждый вызов модуля направляется следующему InterpreterModule, а после последнего выполнение возвращается к первому.

torch.export.unflatten.unflatten(module, flat_args_adapter=None) [исходный код]

Восстанавливает иерархию модулей в ExportedProgram, создавая модуль с той же иерархией, что и исходный модуль eager-режима. Это может быть полезно, если вы хотите использовать torch.export с другой системой, которая ожидает иерархию модулей, а не плоский граф, обычно создаваемый torch.export.

Примечание

Аргументы args/kwargs восстановленных модулей не обязательно будут соответствовать исходному модулю в eager-режиме, поэтому замена модуля (например, self.submod = new_mod) может не сработать. Если нужно заменить модуль, необходимо задать параметр preserve_module_call_signature в torch.export.export().

Параметры:

module (ExportedProgram) – ExportedProgram, для которой нужно восстановить иерархию.
flat_args_adapter (Optional[FlatArgsAdapter]) – Адаптирует плоские аргументы, если входной TreeSpec не соответствует TreeSpec экспортированного модуля.

Возвращает:

Экземпляр UnflattenedModule с той же иерархией модулей, что и у исходного модуля eager-режима до экспорта.

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

UnflattenedModule

torch.export.register_dataclass(cls, *, serialized_type_name=None) [исходный код]

Регистрирует dataclass как допустимый тип входных и выходных данных для torch.export.export().

Параметры:

cls (type[Any]) – тип dataclass для регистрации
serialized_type_name (str | None) – Сериализованное имя dataclass. Оно
this (обязательно, если нужно сериализовать содержащий) –
dataclass. –

Пример:

import torch
from dataclasses import dataclass


@dataclass
class InputDataClass:
    feature: torch.Tensor
    bias: int


@dataclass
class OutputDataClass:
    res: torch.Tensor


torch.export.register_dataclass(InputDataClass)
torch.export.register_dataclass(OutputDataClass)


class Mod(torch.nn.Module):
    def forward(self, x: InputDataClass) -> OutputDataClass:
        res = x.feature + x.bias
        return OutputDataClass(res=res)


ep = torch.export.export(Mod(), (InputDataClass(torch.ones(2, 2), 1),))
print(ep)

class torch.export.decomp_utils.CustomDecompTable [исходный код]

Это специальный пользовательский словарь для обработки decomp_table при экспорте. Он необходим потому, что в новой модели для сохранения оператора его можно только удалить из таблицы декомпозиций. Это проблематично для пользовательских операторов, поскольку неизвестно, когда именно такой оператор будет загружен в диспетчер. Поэтому операции с пользовательскими операторами нужно записывать и материализовать только тогда, когда это действительно необходимо (при запуске этапа декомпозиции).

Соблюдаются следующие инварианты:

Все декомпозиции aten загружаются при инициализации.
При любом чтении таблицы материализуются ВСЕ операторы, чтобы увеличить вероятность того, что диспетчер обнаружит пользовательский оператор.
При операции записи материализация не всегда требуется.
Окончательная загрузка выполняется во время экспорта, непосредственно перед вызовом run_decompositions().

copy() [исходный код]

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

items() [исходный код]

keys() [исходный код]

materialize() [исходный код]

Тип возвращаемого значения:: dict[OperatorBase, Callable]

pop(*args) [исходный код]

update(other_dict) [исходный код]

torch.export.passes.move_to_device_pass(ep, location) [исходный код]

Переносит экспортированную программу на указанное устройство.

Параметры:

ep (ExportedProgram) – Экспортированная программа для переноса.
location (Union[torch.device, str, Dict[str, str]]) – Устройство, на которое нужно перенести экспортированную программу. Если указана строка, она интерпретируется как имя устройства. Если указан словарь, он интерпретируется как сопоставление исходных устройств с целевыми.

Возвращает:

Перенесённая экспортированная программа.

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

ExportedProgram

class torch.export.pt2_archive.PT2ArchiveReader(archive_path_or_buffer)

Менеджер контекста для чтения архива PT2.

archive_version() [источник]

Получить версию архива.

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

get_file_names() [источник]

Получить имена файлов в архиве.

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

read_bytes(name) [источник]

Прочитать объект bytes из архива. name: исходный файл внутри архива.

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

read_string(name) [источник]

Прочитать строковый объект из архива. name: исходный файл внутри архива.

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

class torch.export.pt2_archive.PT2ArchiveWriter(archive_path_or_buffer)

Менеджер контекста для записи архива PT2.

close() [источник]: Закрыть архив.

count_prefix(prefix) [источник]

Подсчитать количество записей, начинающихся с заданного префикса.

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

has_record(name) [источник]

Проверить, существует ли запись в архиве.

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

write_bytes(name, data) [источник]: Записать объект bytes в архив. name: целевой файл внутри архива. data: объект bytes для записи.

write_file(name, file_path) [источник]: Скопировать файл в архив. name: целевой файл внутри архива. file_path: исходный файл на диске.

write_folder(archive_dir, folder_dir) [источник]: Скопировать папку в архив. archive_dir: целевая папка внутри архива. folder_dir: исходная папка на диске.

write_string(name, data) [источник]: Записать строковый объект в архив. name: целевой файл внутри архива. data: строковый объект для записи.

torch.export.pt2_archive.is_pt2_package(serialized_model) [источник]

Проверить, является ли сериализованная модель пакетом PT2 Archive.

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

class torch.export.exported_program.ModuleCallEntry(fqn: str, signature: torch.export.exported_program.ModuleCallSignature | None = None) [источник]

class torch.export.exported_program.ModuleCallSignature(inputs: list[torch.export.graph_signature.TensorArgument | torch.export.graph_signature.SymIntArgument | torch.export.graph_signature.SymFloatArgument | torch.export.graph_signature.SymBoolArgument | torch.export.graph_signature.ConstantArgument | torch.export.graph_signature.CustomObjArgument | torch.export.graph_signature.TokenArgument], outputs: list[torch.export.graph_signature.TensorArgument | torch.export.graph_signature.SymIntArgument | torch.export.graph_signature.SymFloatArgument | torch.export.graph_signature.SymBoolArgument | torch.export.graph_signature.ConstantArgument | torch.export.graph_signature.CustomObjArgument | torch.export.graph_signature.TokenArgument], in_spec: torch.utils._pytree.TreeSpec, out_spec: torch.utils._pytree.TreeSpec, forward_arg_names: list[str] | None = None) [источник]

torch.export.exported_program.default_decompositions() [источник]

Это таблица декомпозиций по умолчанию, содержащая декомпозиции всех операторов ATen до базового набора операторов ATen. Используйте этот API вместе с run_decompositions()

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

class torch.export.custom_obj.ScriptObjectMeta(constant_name, class_fqn) [источник]: Метаданные, хранящиеся в узлах, представляющих объекты ScriptObject.

class torch.export.graph_signature.ConstantArgument(name: str, value: int | float | bool | str | None) [источник]

name: str

value: int | float | bool | str | None

class torch.export.graph_signature.CustomObjArgument(name: str, class_fqn: str, fake_val: torch._library.fake_class_registry.FakeScriptObject | None = None) [источник]

class_fqn: str

fake_val: FakeScriptObject | None = None

name: str

class torch.export.graph_signature.ExportBackwardSignature(gradients_to_parameters: dict[str, str], gradients_to_user_inputs: dict[str, str], loss_output: str) [источник]

gradients_to_parameters: dict[str, str]

gradients_to_user_inputs: dict[str, str]

loss_output: str

class torch.export.graph_signature.ExportGraphSignature(input_specs, output_specs) [источник]

ExportGraphSignature описывает сигнатуру входов и выходов графа экспорта — fx.Graph с более строгими гарантиями инвариантности.

Граф экспорта является функциональным и не обращается к «состояниям», таким как параметры или буферы, внутри графа с помощью узлов getattr. Вместо этого export() гарантирует, что параметры, буферы и константные тензоры выносятся из графа и становятся входами. Аналогичным образом, любые изменения буферов также не включаются в граф: обновлённые значения изменённых буферов моделируются как дополнительные выходы графа экспорта.

Порядок всех входов и выходов:

Inputs = [*parameters_buffers_constant_tensors, *flattened_user_inputs]
Outputs = [*mutated_inputs, *flattened_user_outputs]

Например, если экспортировать следующий модуль:

class CustomModule(nn.Module):
    def __init__(self) -> None:
        super(CustomModule, self).__init__()

        # Define a parameter
        self.my_parameter = nn.Parameter(torch.tensor(2.0))

        # Define two buffers
        self.register_buffer("my_buffer1", torch.tensor(3.0))
        self.register_buffer("my_buffer2", torch.tensor(4.0))

    def forward(self, x1, x2):
        # Use the parameter, buffers, and both inputs in the forward method
        output = (
            x1 + self.my_parameter
        ) * self.my_buffer1 + x2 * self.my_buffer2

        # Mutate one of the buffers (e.g., increment it by 1)
        self.my_buffer2.add_(1.0)  # In-place addition

        return output


mod = CustomModule()
ep = torch.export.export(mod, (torch.tensor(1.0), torch.tensor(2.0)))

Полученный граф является нефункциональным:

graph():
    %p_my_parameter : [num_users=1] = placeholder[target=p_my_parameter]
    %b_my_buffer1 : [num_users=1] = placeholder[target=b_my_buffer1]
    %b_my_buffer2 : [num_users=2] = placeholder[target=b_my_buffer2]
    %x1 : [num_users=1] = placeholder[target=x1]
    %x2 : [num_users=1] = placeholder[target=x2]
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x1, %p_my_parameter), kwargs = {})
    %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, %b_my_buffer1), kwargs = {})
    %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%x2, %b_my_buffer2), kwargs = {})
    %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %mul_1), kwargs = {})
    %add_ : [num_users=0] = call_function[target=torch.ops.aten.add_.Tensor](args = (%b_my_buffer2, 1.0), kwargs = {})
    return (add_1,)

Результирующая ExportGraphSignature нефункционального графа будет выглядеть так:

# inputs
p_my_parameter: PARAMETER target='my_parameter'
b_my_buffer1: BUFFER target='my_buffer1' persistent=True
b_my_buffer2: BUFFER target='my_buffer2' persistent=True
x1: USER_INPUT
x2: USER_INPUT

# outputs
add_1: USER_OUTPUT

Чтобы получить функциональный граф, можно использовать run_decompositions():

mod = CustomModule()
ep = torch.export.export(mod, (torch.tensor(1.0), torch.tensor(2.0)))
ep = ep.run_decompositions()

Полученный граф является функциональным:

graph():
    %p_my_parameter : [num_users=1] = placeholder[target=p_my_parameter]
    %b_my_buffer1 : [num_users=1] = placeholder[target=b_my_buffer1]
    %b_my_buffer2 : [num_users=2] = placeholder[target=b_my_buffer2]
    %x1 : [num_users=1] = placeholder[target=x1]
    %x2 : [num_users=1] = placeholder[target=x2]
    %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%x1, %p_my_parameter), kwargs = {})
    %mul : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%add, %b_my_buffer1), kwargs = {})
    %mul_1 : [num_users=1] = call_function[target=torch.ops.aten.mul.Tensor](args = (%x2, %b_my_buffer2), kwargs = {})
    %add_1 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%mul, %mul_1), kwargs = {})
    %add_2 : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%b_my_buffer2, 1.0), kwargs = {})
    return (add_2, add_1)

Результирующая ExportGraphSignature функционального графа будет выглядеть так:

# inputs
p_my_parameter: PARAMETER target='my_parameter'
b_my_buffer1: BUFFER target='my_buffer1' persistent=True
b_my_buffer2: BUFFER target='my_buffer2' persistent=True
x1: USER_INPUT
x2: USER_INPUT

# outputs
add_2: BUFFER_MUTATION target='my_buffer2'
add_1: USER_OUTPUT

property assertion_dep_token: Mapping[int, str] | None

property backward_signature: ExportBackwardSignature | None

property buffers: Collection[str]

property buffers_to_mutate: Mapping[str, str]

get_replace_hook(replace_inputs=False) [источник]

input_specs: list[InputSpec]

property input_tokens: Collection[str]

property inputs_to_buffers: Mapping[str, str]

property inputs_to_lifted_custom_objs: Mapping[str, str]

property inputs_to_lifted_tensor_constants: Mapping[str, str]

property inputs_to_parameters: Mapping[str, str]

property lifted_custom_objs: Collection[str]

property lifted_tensor_constants: Collection[str]

property non_persistent_buffers: Collection[str]

output_specs: list[OutputSpec]

property output_tokens: Collection[str]

property parameters: Collection[str]

property parameters_to_mutate: Mapping[str, str]

replace_all_uses(old, new) [источник]: Заменить в сигнатуре все вхождения старого имени новым именем.

property user_inputs: Collection[int | float | bool | str | None]

property user_inputs_to_mutate: Mapping[str, str]

property user_outputs: Collection[int | float | bool | str | None]

class torch.export.graph_signature.InputKind(value) [источник]

Перечисление.

BUFFER = 3

CONSTANT_TENSOR = 4

CUSTOM_OBJ = 5

PARAMETER = 2

TOKEN = 6

USER_INPUT = 1

class torch.export.graph_signature.InputSpec(kind: torch.export.graph_signature.InputKind, arg: torch.export.graph_signature.TensorArgument | torch.export.graph_signature.SymIntArgument | torch.export.graph_signature.SymFloatArgument | torch.export.graph_signature.SymBoolArgument | torch.export.graph_signature.ConstantArgument | torch.export.graph_signature.CustomObjArgument | torch.export.graph_signature.TokenArgument, target: str | None, persistent: bool | None = None) [источник]

arg: TensorArgument | SymIntArgument | SymFloatArgument | SymBoolArgument | ConstantArgument | CustomObjArgument | TokenArgument

kind: InputKind

persistent: bool | None = None

target: str | None

class torch.export.graph_signature.OutputKind(value) [источник]

Перечисление.

BUFFER_MUTATION = 3

GRADIENT_TO_PARAMETER = 5

GRADIENT_TO_USER_INPUT = 6

LOSS_OUTPUT = 2

PARAMETER_MUTATION = 4

TOKEN = 8

USER_INPUT_MUTATION = 7

USER_OUTPUT = 1

class torch.export.graph_signature.OutputSpec(kind: torch.export.graph_signature.OutputKind, arg: torch.export.graph_signature.TensorArgument | torch.export.graph_signature.SymIntArgument | torch.export.graph_signature.SymFloatArgument | torch.export.graph_signature.SymBoolArgument | torch.export.graph_signature.ConstantArgument | torch.export.graph_signature.CustomObjArgument | torch.export.graph_signature.TokenArgument, target: str | None) [источник]

arg: TensorArgument | SymIntArgument | SymFloatArgument | SymBoolArgument | ConstantArgument | CustomObjArgument | TokenArgument

kind: OutputKind

target: str | None

class torch.export.graph_signature.SymBoolArgument(name: str) [источник]

name: str

class torch.export.graph_signature.SymFloatArgument(name: str) [источник]

name: str

class torch.export.graph_signature.SymIntArgument(name: str) [источник]

name: str

class torch.export.graph_signature.TensorArgument(name: str) [источник]

name: str

class torch.export.graph_signature.TokenArgument(name: str) [источник]

name: str

© 2026, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://docs.pytorch.org/docs/2.14/user_guide/torch_compiler/export/api_reference.html