torch.export

Обзор

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

import torch
from torch.export import export

def f(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b

example_args = (torch.randn(10, 10), torch.randn(10, 10))

exported_program: torch.export.ExportedProgram = export(
    f, args=example_args
)
print(exported_program)

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: f32[10, 10], arg1_1: f32[10, 10]):
            # code: a = torch.sin(x)
            sin: f32[10, 10] = torch.ops.aten.sin.default(arg0_1);

            # code: b = torch.cos(y)
            cos: f32[10, 10] = torch.ops.aten.cos.default(arg1_1);

            # code: return a + b
            add: f32[10, 10] = torch.ops.aten.add.Tensor(sin, cos);
            return (add,)

    Graph signature: ExportGraphSignature(
        parameters=[],
        buffers=[],
        user_inputs=['arg0_1', 'arg1_1'],
        user_outputs=['add'],
        inputs_to_parameters={},
        inputs_to_buffers={},
        buffers_to_mutate={},
        backward_signature=None,
        assertion_dep_token=None,
    )
    Range constraints: {}
    Equality constraints: []

torch.export генерирует чистую промежуточную репрезентацию (IR) с соблюдением следующих инвариантов. Более подробные сведения о IR можно найти здесь (скоро).

• Согласованность: Гарантируется, что это звуковая репрезентация исходной программы и сохраняются те же соглашения о вызовах, что и в исходной программе.
• Нормализация: В графе нет семантики Python. Подмодули из исходных программ инлайнят для формирования одного полностью уплощенного вычислительного графа.
• Определенный набор операторов: Сгенерированный граф содержит только небольшой набор определенных операторов ядра ATen IR и зарегистрированных пользовательских операторов.
• Свойства графа: Граф является чисто функциональным, то есть не содержит операций с побочными эффектами, таких как изменения или алиасинг. Он не изменяет промежуточные значения, параметры или буферы.
• Метаданные: Граф содержит метаданные, полученные во время прослеживания, такие как стек вызовов из кода пользователя.

Внутренне, torch.export использует следующие новейшие технологии:

• TorchDynamo (torch._dynamo) — это внутренний API, который использует функцию CPython, называемую API оценки фреймов, для безопасного прослеживания графов PyTorch. Это обеспечивает значительно улучшенное прослеживание графа, с гораздо меньшим количеством переписываний, необходимых для полного прослеживания кода PyTorch.
• AOT Autograd предоставляет функционализированный граф PyTorch и гарантирует, что граф декомпозирован/опущен до небольшого набора операторов ядра ATen.
• Torch FX (torch.fx) — это основное представление графа, позволяющее гибкие преобразования на основе Python.

Существующие фреймворки

torch.compile() также использует ту же стопку PT2, что и torch.export, но немного отличается:

• JIT vs. AOT: torch.compile() — это компилятор JIT, который не предназначен для создания скомпилированных артефактов за пределами развертывания.
• Частичное vs. Полное захват графа: Когда torch.compile() сталкивается с непрослеживаемой частью модели, он «разрывает граф» и возвращается к выполнению программы в режиме выполнения Python eager. В сравнении, torch.export нацелен на получение полного представления графа модели PyTorch, поэтому он будет выдавать ошибку, когда будет достигнута непрослеживаемая часть. Поскольку torch.export генерирует полный граф, отделенный от любых функций или среды выполнения Python, этот граф можно сохранить, загрузить и выполнить в разных средах и языках.
• Компромисс в плане удобства использования: Так как torch.compile() может переключаться на среду выполнения Python, когда встречает непрослеживаемую часть, он гораздо более гибкий. torch.export вместо этого потребует от пользователей предоставления дополнительной информации или переписывания кода для обеспечения возможности прослеживания.

По сравнению с torch.fx.symbolic_trace(), torch.export прослеживает с помощью TorchDynamo, который работает на уровне байткода Python, позволяя ему прослеживать произвольные конструкции Python, не ограниченные тем, что поддерживает перегрузка операторов Python. Кроме того, torch.export отслеживает подробные метаданные тензоров, чтобы условные операторы на таких параметрах, как форма тензора, не мешали прослеживанию. В целом, torch.export должен работать с большей частью пользовательских программ и генерировать графы более низкого уровня (на уровне операторов torch.ops.aten). Обратите внимание, что пользователи все еще могут использовать torch.fx.symbolic_trace() как предварительную стадию перед torch.export.

По сравнению с torch.jit.script(), torch.export не захватывает поток управления Python или структуры данных, но он поддерживает больше функций языка Python, чем TorchScript (так как проще иметь полное покрытие байткодов Python). Полученные графы проще и содержат только прямой поток управления (за исключением явных операторов потока управления).

По сравнению с torch.jit.trace(), torch.export согласован: он может проследить код, который выполняет целочисленные вычисления на размерах, и записывает все необходимые дополнительные условия, чтобы показать, что конкретный прослеженный граф действителен для других входных данных.

Экспорт модели PyTorch

Пример

Основным входным точкой является torch.export.export(), который принимает вызываемый объект (torch.nn.Module, функцию или метод) и примеры входных данных, и сохраняет вычислительный граф в виде torch.export.ExportedProgram. Пример:

import torch
from torch.export import export

# Simple module for demonstration
class M(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = torch.nn.Conv2d(
            in_channels=3, out_channels=16, kernel_size=3, padding=1
        )
        self.relu = torch.nn.ReLU()
        self.maxpool = torch.nn.MaxPool2d(kernel_size=3)

    def forward(self, x: torch.Tensor, *, constant=None) -> torch.Tensor:
        a = self.conv(x)
        a.add_(constant)
        return self.maxpool(self.relu(a))

example_args = (torch.randn(1, 3, 256, 256),)
example_kwargs = {"constant": torch.ones(1, 16, 256, 256)}

exported_program: torch.export.ExportedProgram = export(
    M(), args=example_args, kwargs=example_kwargs
)
print(exported_program)

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: f32[16, 3, 3, 3], arg1_1: f32[16], arg2_1: f32[1, 3, 256, 256], arg3_1: f32[1, 16, 256, 256]):

            # code: a = self.conv(x)
            convolution: f32[1, 16, 256, 256] = torch.ops.aten.convolution.default(
                arg2_1, arg0_1, arg1_1, [1, 1], [1, 1], [1, 1], False, [0, 0], 1
            );

            # code: a.add_(constant)
            add: f32[1, 16, 256, 256] = torch.ops.aten.add.Tensor(convolution, arg3_1);

            # code: return self.maxpool(self.relu(a))
            relu: f32[1, 16, 256, 256] = torch.ops.aten.relu.default(add);
            max_pool2d_with_indices = torch.ops.aten.max_pool2d_with_indices.default(
                relu, [3, 3], [3, 3]
            );
            getitem: f32[1, 16, 85, 85] = max_pool2d_with_indices[0];
            return (getitem,)

    Graph signature: ExportGraphSignature(
        parameters=['L__self___conv.weight', 'L__self___conv.bias'],
        buffers=[],
        user_inputs=['arg2_1', 'arg3_1'],
        user_outputs=['getitem'],
        inputs_to_parameters={
            'arg0_1': 'L__self___conv.weight',
            'arg1_1': 'L__self___conv.bias',
        },
        inputs_to_buffers={},
        buffers_to_mutate={},
        backward_signature=None,
        assertion_dep_token=None,
    )
    Range constraints: {}
    Equality constraints: []

Просматривая ExportedProgram, мы можем отметить следующее:

• Граф torch.fx.Graph содержит вычислительный граф исходной программы, а также записи исходного кода для удобства отладки.
• Граф содержит только операторы torch.ops.aten , обнаруженные в наборе операторов ядра ATen IR и пользовательских операторов, и является полностью функциональным без операторов типа in-place, таких как torch.add_.
• Параметры (веса и смещения conv) подняты в качестве входных данных для графа, что приводит к отсутствию узлов get_attr в графе, которые ранее существовали в результате torch.fx.symbolic_trace().
• torch.export.ExportGraphSignature моделирует входные и выходные сигнатуры, а также определяет, какие входные данные являются параметрами.
• Записываются форма и тип данных тензоров, создаваемых каждым узлом в графе. Например, узел convolution приведет к тензору с типом данных torch.float32 и формой (1, 16, 256, 256).

Выражение динамики

По умолчанию torch.export проследит программу, предполагая, что все формы входных данных являются **статическими**, и специализирует экспортированную программу на этих размерах. Однако некоторые размеры, такие как размер пакетного измерения, могут быть динамическими и изменяться от выполнения к выполнению. Такие размеры должны быть отмечены как динамические с помощью API torch.export.dynamic_dim() и переданы в torch.export.export() в качестве аргумента constraints. Пример:

import torch
from torch.export import export, dynamic_dim

class M(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.branch1 = torch.nn.Sequential(
            torch.nn.Linear(64, 32), torch.nn.ReLU()
        )
        self.branch2 = torch.nn.Sequential(
            torch.nn.Linear(128, 64), torch.nn.ReLU()
        )
        self.buffer = torch.ones(32)

    def forward(self, x1, x2):
        out1 = self.branch1(x1)
        out2 = self.branch2(x2)
        return (out1 + self.buffer, out2)

example_args = (torch.randn(32, 64), torch.randn(32, 128))
constraints = [
    # First dimension of each input is a dynamic batch size
    dynamic_dim(example_args[0], 0),
    dynamic_dim(example_args[1], 0),
    # The dynamic batch size between the inputs are equal
    dynamic_dim(example_args[0], 0) == dynamic_dim(example_args[1], 0),
]

exported_program: torch.export.ExportedProgram = export(
  M(), args=example_args, constraints=constraints
)
print(exported_program)

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: f32[32, 64], arg1_1: f32[32], arg2_1: f32[64, 128], arg3_1: f32[64], arg4_1: f32[32], arg5_1: f32[s0, 64], arg6_1: f32[s0, 128]):

            # code: out1 = self.branch1(x1)
            permute: f32[64, 32] = torch.ops.aten.permute.default(arg0_1, [1, 0]);
            addmm: f32[s0, 32] = torch.ops.aten.addmm.default(arg1_1, arg5_1, permute);
            relu: f32[s0, 32] = torch.ops.aten.relu.default(addmm);

            # code: out2 = self.branch2(x2)
            permute_1: f32[128, 64] = torch.ops.aten.permute.default(arg2_1, [1, 0]);
            addmm_1: f32[s0, 64] = torch.ops.aten.addmm.default(arg3_1, arg6_1, permute_1);
            relu_1: f32[s0, 64] = torch.ops.aten.relu.default(addmm_1);  addmm_1 = None

            # code: return (out1 + self.buffer, out2)
            add: f32[s0, 32] = torch.ops.aten.add.Tensor(relu, arg4_1);
            return (add, relu_1)

    Graph signature: ExportGraphSignature(
        parameters=[
            'branch1.0.weight',
            'branch1.0.bias',
            'branch2.0.weight',
            'branch2.0.bias',
        ],
        buffers=['L__self___buffer'],
        user_inputs=['arg5_1', 'arg6_1'],
        user_outputs=['add', 'relu_1'],
        inputs_to_parameters={
            'arg0_1': 'branch1.0.weight',
            'arg1_1': 'branch1.0.bias',
            'arg2_1': 'branch2.0.weight',
            'arg3_1': 'branch2.0.bias',
        },
        inputs_to_buffers={'arg4_1': 'L__self___buffer'},
        buffers_to_mutate={},
        backward_signature=None,
        assertion_dep_token=None,
    )
    Range constraints: {s0: RangeConstraint(min_val=2, max_val=9223372036854775806)}
    Equality constraints: [(InputDim(input_name='arg5_1', dim=0), InputDim(input_name='arg6_1', dim=0))]

Некоторые дополнительные моменты:

• С помощью API torch.export.dynamic_dim() мы указали, что первый размер каждого входного значения является динамическим. Рассматривая входы arg5_1 и arg6_1, у них есть символьная форма (s0, 64) и (s0, 128) вместо тензоров с формой (32, 64) и (32, 128), которые мы передали в качестве примеров входных данных. s0 — это символ, обозначающий, что этот размер может принимать ряд значений.
• exported_program.range_constraints описывает диапазоны каждого символа, появляющегося в графе. В данном случае мы видим, что s0 имеет диапазон [2, ∞]. По техническим причинам, которые трудно объяснить здесь, они предполагаются не равными 0 или 1. Это не ошибка и не обязательно означает, что экспортированная программа не будет работать для размеров 0 или 1. См. Проблема специализации 0/1 для углубленного обсуждения этой темы.
• exported_program.equality_constraints описывает, какие размеры должны быть равны. Поскольку мы указали в ограничениях, что первый размер каждого аргумента эквивалентен (dynamic_dim(example_args[0], 0) == dynamic_dim(example_args[1], 0)), в ограничениях равенства мы видим кортеж, указывающий, что размер arg5_1 0 равен размеру arg6_1 0.

Сериализация

Для сохранения ExportedProgram, пользователи могут использовать API torch.export.save() и torch.export.load(). Используйте соглашение о сохранении ExportedProgram с расширением файла .pt2.

Пример:

import torch
import io

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

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

torch.export.save(exported_program, 'exported_program.pt2')
saved_exported_program = torch.export.load('exported_program.pt2')

Специализация

Формы входных данных

Как упоминалось ранее, по умолчанию torch.export проследит программу, специализируясь на формах тензоров входных данных, если не указан размер как динамический через API torch.export.dynamic_dim(). Это означает, что если существует поток управления, зависящий от формы, torch.export будет специализироваться на ветви, которая выполняется с заданными входными данными. Например:

import torch
from torch.export import export

def fn(x):
    if x.shape[0] > 5:
        return x + 1
    else:
        return x - 1

example_inputs = (torch.rand(10, 2),)
exported_program = export(fn, example_inputs)
print(exported_program)

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: f32[10, 2]):
            add: f32[10, 2] = torch.ops.aten.add.Tensor(arg0_1, 1);
            return (add,)

Условие (x.shape[0] > 5) не появляется в ExportedProgram, потому что входные примеры имеют статическую форму (10, 2). Поскольку torch.export специализируется на статических формах входных данных, ветвь else (x - 1) никогда не будет достигнута. Чтобы сохранить динамическое ветвление, основанное на форме тензора в отслеживаемом графе, необходимо использовать torch.export.dynamic_dim(), чтобы указать, что размер входного тензора (x.shape[0]) динамический, и исходный код необходимо переписать.

Входы, не являющиеся тензорами

torch.export также специализирует отслеживаемый граф на основе значений входов, которые не являются torch.Tensor, таких как int, float, bool, и str. Однако, мы, вероятно, изменим это в ближайшем будущем, чтобы не специализироваться на входных данных примитивных типов.

Например:

import torch
from torch.export import export

def fn(x: torch.Tensor, const: int, times: int):
    for i in range(times):
        x = x + const
    return x

example_inputs = (torch.rand(2, 2), 1, 3)
exported_program = export(fn, example_inputs)
print(exported_program)

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: f32[2, 2], arg1_1, arg2_1):
            add: f32[2, 2] = torch.ops.aten.add.Tensor(arg0_1, 1);
            add_1: f32[2, 2] = torch.ops.aten.add.Tensor(add, 1);
            add_2: f32[2, 2] = torch.ops.aten.add.Tensor(add_1, 1);
            return (add_2,)

Поскольку целые числа специализированы, операции torch.ops.aten.add.Tensor все вычисляются с встроенной константой 1, а не с arg1_1. Кроме того, итератор times используемый в цикле for также «встраивается» в граф через 3 повторяющиеся вызовы torch.ops.aten.add.Tensor, и вход arg2_1 никогда не используется.

Ограничения torch.export

Разрывы графа

Поскольку torch.export — это одноразовый процесс захвата вычислительного графа из программы PyTorch, он может в конечном итоге столкнуться с неотслеживаемыми частями программ, поскольку практически невозможно отследить все функции PyTorch и Python. В случае torch.compile, недопустимая операция вызовет «разрыв графа», и недопустимая операция будет выполнена с использованием стандартного Python-вычисления. В отличие от этого, torch.export потребует от пользователей предоставления дополнительной информации или переписывания частей кода, чтобы сделать его отслеживаемым. Поскольку отслеживание основано на TorchDynamo, который оценивает на уровне байткода Python, потребуется значительно меньше переписываний по сравнению с предыдущими системами отслеживания.

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

Управляющие конструкции, зависящие от данных/формы

Разрывы графа также могут возникать при управляющих конструкциях, зависящих от данных (if x.shape[0] > 2) , когда формы не специализируются, так как компилятор отслеживания не может справиться с этим без генерации кода для комбинаториально взрывающегося количества путей. В таких случаях пользователям необходимо переписать свой код, используя специальные операторы управляющих конструкций (скоро появятся!).

Доступ, зависящий от данных

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

Отсутствующие мета-ядра для операторов

При отслеживании для всех операторов требуется реализация META (или «мета-ядро»). Это используется для определения форм входных/выходных данных для этого оператора.

Обратите внимание, что официальный API для регистрации пользовательских мета-ядер для пользовательских операций в настоящее время находится в разработке. Пока окончательный API уточняется, вы можете обратиться к документации здесь.

В случае, если ваша модель использует оператор ATen, для которого еще нет реализации мета-ядра, пожалуйста, создайте вопрос.

Дополнительная информация

Дополнительные ссылки для пользователей Export

• Создание преобразований графа в ATen IR
• IR
• ExportDB

Глубокое погружение для разработчиков PyTorch

• Глубокое погружение в TorchDynamo
• Динамические формы
• Фиктивный тензор

Справочник API

torch.export.export(f, args, kwargs=None, *, constraints=None) [source]

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

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

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

Существует 2 типа предположений, сделанных во время прослеживания:

• Формы (а не значения) входных тензоров.
• Диапазоны (нижняя и верхняя границы) значений, извлечённых из промежуточных тензоров с помощью .item() или прямого индексирования.

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

• Предположения о статических формах входных тензоров автоматически проверяются без дополнительных усилий.
• Предположения о динамических формах входных тензоров требуют явного Input Constraint, созданного с помощью API dynamic_dim()
• Предположения о диапазоне промежуточных значений требуют явного Inline Constraint, созданного с помощью API constrain_as_size() и API constraint_as_value().

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

def specify_constraints(x):
    return [
        # x:
        dynamic_dim(x, 0) <= 5,
    ]

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

Parameters

• f (Callable) – Функция, которую нужно проследить.
• args (Tuple[Any, ...]) – Пример позиционных входных данных.
• kwargs (Optional[Dict[str, Any]]) – Необязательные примеры ключевых входных данных.
• constraints (Optional[List[Constraint]]) – Необязательный список ограничений на динамические аргументы, которые задают их возможный диапазон форм. По умолчанию формы входных torch.Tensors предполагаются статическими. Если ожидается, что входной torch.Tensor будет иметь динамические формы, используйте dynamic_dim() для определения объектов Constraint, которые определяют динамику и возможный диапазон форм. См. документацию dynamic_dim() для примеров использования.

Returns

Экземпляр ExportedProgram с прослеженной функцией.

Return type

ExportedProgram

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

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

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

torch.export.dynamic_dim(t, index) [source]

dynamic_dim() создаёт объект Constraint, описывающий динамизм измерения index тензора t. Объекты Constraint должны быть переданы в аргумент constraints функции export().

Parameters

• t (torch.Tensor) – Пример входного тензора, имеющего динамический размер измерения(ей)
• index (int) – Индекс динамического измерения

Returns

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

В частности, dynamic_dim() может использоваться для выражения следующих типов динамизма.

• Размер измерения динамичен и не ограничен:

  t0 = torch.rand(2, 3)
  t1 = torch.rand(3, 4)
  
  # First dimension of t0 can be dynamic size rather than always being static size 2
  constraints = [dynamic_dim(t0, 0)]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• Размер измерения динамичен с нижней границей:

  t0 = torch.rand(10, 3)
  t1 = torch.rand(3, 4)
  
  # First dimension of t0 can be dynamic size with a lower bound of 5 (inclusive)
  # Second dimension of t1 can be dynamic size with a lower bound of 2 (exclusive)
  constraints = [
      dynamic_dim(t0, 0) >= 5,
      dynamic_dim(t1, 1) > 2,
  ]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• Размер измерения динамичен с верхней границей:

  t0 = torch.rand(10, 3)
  t1 = torch.rand(3, 4)
  
  # First dimension of t0 can be dynamic size with a upper bound of 16 (inclusive)
  # Second dimension of t1 can be dynamic size with a upper bound of 8 (exclusive)
  constraints = [
      dynamic_dim(t0, 0) <= 16,
      dynamic_dim(t1, 1) < 8,
  ]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• Размер измерения динамичен и всегда равен размеру другого динамического измерения:

  t0 = torch.rand(10, 3)
  t1 = torch.rand(3, 4)
  
  # Sizes of second dimension of t0 and first dimension are always equal
  constraints = [
      dynamic_dim(t0, 1) == dynamic_dim(t1, 0),
  ]
  ep = export(fn, (t0, t1), constraints=constraints)
  

• Можно комбинировать все вышеперечисленные типы, если они не выражают противоречивые требования.

torch.export.constrain_as_size(symbol, min=None, max=None) [source]

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

Параметры

• symbol – Промежуточное скалярное значение (только int сейчас) для применения ограничения диапазона.
• min (Необязательно[int]) – Минимальное возможное значение данного символа (включительно)
• max (Необязательно[int]) – Максимальное возможное значение данного символа (включительно)

Возвращает

None

Например, следующая программа не может быть отслежена надёжно без использования constrain_as_size() для подсказки export() о диапазонах форм:

def fn(x):
    d = x.max().item()
    return torch.ones(v)

export() выдаст следующую ошибку:

torch._dynamo.exc.Unsupported: guard on data-dependent symbolic int/float

Предполагая, что фактический диапазон d может быть между [3, 10], вы можете добавить вызов constrain_as_size() в исходный код так:

def fn(x):
    d = x.max().item()
    torch.export.constrain_as_size(d, min=3, max=10)
    return torch.ones(d)

С дополнительной подсказкой export() сможет отследить программу правильно, взяв else ветвь, что приведёт к следующей графе:

graph():
    %arg0_1 := placeholder[target=arg0_1]

    # d = x.max().item()
    %max_1 := call_function[target=torch.ops.aten.max.default](args = (%arg0_1,))
    %_local_scalar_dense := call_function[target=torch.ops.aten._local_scalar_dense.default](args = (%max_1,))

    # Asserting 3 <= d <= 10
    %ge := call_function[target=operator.ge](args = (%_local_scalar_dense, 3))
    %scalar_tensor := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%ge,))
    %_assert_async := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor, _local_scalar_dense is outside of inline constraint [3, 10].))
    %le := call_function[target=operator.le](args = (%_local_scalar_dense, 10))
    %scalar_tensor_1 := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%le,))
    %_assert_async_1 := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor_1, _local_scalar_dense is outside of inline constraint [3, 10].))
    %sym_constrain_range_for_size := call_function[target=torch.ops.aten.sym_constrain_range_for_size.default](
        args = (%_local_scalar_dense,), kwargs = {min: 3, max: 10})

    # Constructing new tensor with d
    %full := call_function[target=torch.ops.aten.full.default](
        args = ([%_local_scalar_dense], 1),
        kwargs = {dtype: torch.float32, layout: torch.strided, device: cpu, pin_memory: False})

    ......

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

Если ваш размер предназначен для динамики, НЕ проверяйте, равны ли размеры 0 или 1, эти значения будут НЕЗАМЕЧАТЕЛЬНО сообщать ложь и будут пропущены

torch.export.constrain_as_value(symbol, min=None, max=None) [source]

Подсказка export() о ограничении промежуточного скалярного значения, чтобы последующее ветвление, проверяющее диапазон вышеупомянутого скалярного значения, могло быть надёжно отслежено.

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

(Обратите внимание, что если промежуточное скалярное значение будет использоваться как размер, включая передачу в качестве аргумента размера в фабрику или представление тензора, вызовите constrain_as_size() вместо этого.)

Параметры

• symbol – Промежуточное скалярное значение (только int сейчас) для применения ограничения диапазона.
• min (Необязательно[int]) – Минимальное возможное значение данного символа (включительно)
• max (Необязательно[int]) – Максимальное возможное значение данного символа (включительно)

Возвращает

None

Например, следующая программа не может быть отслежена надёжно:

def fn(x):
    v = x.max().item()
    if v > 1024:
        return x
    else:
        return x * 2

v — это значение, зависящее от данных, которое предполагается иметь диапазон (-беск., +беск.). export() подсказка о том, какую ветвь выбрать, не сможет определить, является ли принятое решение при отслеживании ветвления правильным или нет. Таким образом, export() выдаст следующую ошибку:

torch._dynamo.exc.UserError: Consider annotating your code using
torch.export.constrain_as_size() or torch.export().constrain_as_value() APIs.
It appears that you're trying to get a value out of symbolic int/float whose value
is data-dependent (and thus we do not know the true value.)  The expression we were
trying to evaluate is f0 > 1024 (unhinted: f0 > 1024).

Предполагая, что фактический диапазон v может быть между [10, 200], вы можете добавить вызов constrain_as_value() в исходный код так:

def fn(x):
    v = x.max().item()

    # Give export() a hint
    torch.export.constrain_as_value(v, min=10, max=200)

    if v > 1024:
        return x
    else:
        return x * 2

С дополнительной подсказкой export() сможет отследить программу правильно, взяв else ветвь, что приведёт к следующей графе:

graph():
    %arg0_1 := placeholder[target=arg0_1]

    # v = x.max().item()
    %max_1 := call_function[target=torch.ops.aten.max.default](args = (%arg0_1,))
    %_local_scalar_dense := call_function[target=torch.ops.aten._local_scalar_dense.default](args = (%max_1,))

    # Asserting 10 <= v <= 200
    %ge := call_function[target=operator.ge](args = (%_local_scalar_dense, 10))
    %scalar_tensor := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%ge,))
    %_assert_async := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor, _local_scalar_dense is outside of inline constraint [10, 200].))
    %le := call_function[target=operator.le](args = (%_local_scalar_dense, 200))
    %scalar_tensor_1 := call_function[target=torch.ops.aten.scalar_tensor.default](args = (%le,))
    %_assert_async_1 := call_function[target=torch.ops.aten._assert_async.msg](
        args = (%scalar_tensor_1, _local_scalar_dense is outside of inline constraint [10, 200].))
    %sym_constrain_range := call_function[target=torch.ops.aten.sym_constrain_range.default](
        args = (%_local_scalar_dense,), kwargs = {min: 10, max: 200})

    # Always taking `else` branch to multiply elements `x` by 2 due to hints above
    %mul := call_function[target=torch.ops.aten.mul.Tensor](args = (%arg0_1, 2), kwargs = {})
    return (mul,)

torch.export.save(ep, f, *, extra_files=None, opset_version=None) [source]

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

Разрабатывается, сохранённые файлы могут быть неиспользуемыми в более новых версиях PyTorch.

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

Параметры

• ep (ExportedProgram) – Сохраняемая экспортированная программа.
• f (Union[str, pathlib.Path, io.BytesIO) – Объект, похожий на файл (должен реализовывать write и flush), или строка, содержащая имя файла.
• extra_files (Необязательно[Dict[str, Any]]) – Сопоставление имени файла с содержимым, которое будет сохранено как часть f.
• opset_version (Необязательно[Dict[str, int]]) – Сопоставление имен opset с версией этого opset

Пример:

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'}
torch.export.save(ep, 'exported_program.pt2', extra_files=extra_files)

torch.export.load(f, *, extra_files=None, expected_opset_version=None) [source]

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

Разрабатывается, сохранённые файлы могут быть неиспользуемыми в более новых версиях PyTorch.

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

Параметры

• ep (ExportedProgram) – Сохраняемая экспортированная программа.
• f (Union[str, pathlib.Path, io.BytesIO) – Объект, похожий на файл (должен реализовывать write и flush), или строка, содержащая имя файла.
• extra_files (Необязательно[Dict[str, Any]]) – Дополнительные имена файлов, указанные в этом словаре, будут загружены, и их содержимое будет сохранено в предоставленном словаре.
• expected_opset_version (Необязательно[Dict[str, int]]) – Словарь имен opset и ожидаемых версий opset

Возвращает

Объект 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'])

class torch.export.Constraint(*args, **kwargs) [source]

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

Не создавайте Constraint напрямую, используйте dynamic_dim() вместо этого.

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

class torch.export.ExportedProgram(root, graph, graph_signature, call_spec, state_dict, range_constraints, equality_constraints, module_call_graph, example_inputs=None) [source]

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

Вы можете вызвать ExportedProgram так же, как и исходную вызываемую функцию, прослеженную export() с тем же соглашением вызова.

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

module() [source]

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

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

Модуль

class torch.export.ExportBackwardSignature(gradients_to_parameters: Dict[str, str], gradients_to_user_inputs: Dict[str, str], loss_output: str) [source]

class torch.export.ExportGraphSignature(parameters, buffers, user_inputs, user_outputs, inputs_to_parameters, inputs_to_buffers, buffers_to_mutate, backward_signature, assertion_dep_token=None) [source]

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

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

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

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

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

class CustomModule(nn.Module):
    def __init__(self):
        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

Получившийся граф будет:

graph():
    %arg0_1 := placeholder[target=arg0_1]
    %arg1_1 := placeholder[target=arg1_1]
    %arg2_1 := placeholder[target=arg2_1]
    %arg3_1 := placeholder[target=arg3_1]
    %arg4_1 := placeholder[target=arg4_1]
    %add_tensor := call_function[target=torch.ops.aten.add.Tensor](args = (%arg3_1, %arg0_1), kwargs = {})
    %mul_tensor := call_function[target=torch.ops.aten.mul.Tensor](args = (%add_tensor, %arg1_1), kwargs = {})
    %mul_tensor_1 := call_function[target=torch.ops.aten.mul.Tensor](args = (%arg4_1, %arg2_1), kwargs = {})
    %add_tensor_1 := call_function[target=torch.ops.aten.add.Tensor](args = (%mul_tensor, %mul_tensor_1), kwargs = {})
    %add_tensor_2 := call_function[target=torch.ops.aten.add.Tensor](args = (%arg2_1, 1.0), kwargs = {})
    return (add_tensor_2, add_tensor_1)

Получившийся ExportGraphSignature будет:

ExportGraphSignature(
    # Indicates that there is one parameter named `my_parameter`
    parameters=['L__self___my_parameter'],

    # Indicates that there are two buffers, `my_buffer1` and `my_buffer2`
    buffers=['L__self___my_buffer1', 'L__self___my_buffer2'],

    # Indicates that the nodes `arg3_1` and `arg4_1` in produced graph map to
    # original user inputs, ie. x1 and x2
    user_inputs=['arg3_1', 'arg4_1'],

    # Indicates that the node `add_tensor_1` maps to output of original program
    user_outputs=['add_tensor_1'],

    # Indicates that there is one parameter (self.my_parameter) captured,
    # its name is now mangled to be `L__self___my_parameter`, which is now
    # represented by node `arg0_1` in the graph.
    inputs_to_parameters={'arg0_1': 'L__self___my_parameter'},

    # Indicates that there are two buffers (self.my_buffer1, self.my_buffer2) captured,
    # their name are now mangled to be `L__self___my_my_buffer1` and `L__self___my_buffer2`.
    # They are now represented by nodes `arg1_1` and `arg2_1` in the graph.
    inputs_to_buffers={'arg1_1': 'L__self___my_buffer1', 'arg2_1': 'L__self___my_buffer2'},

    # Indicates that one buffer named `L__self___my_buffer2` is mutated during execution,
    # its new value is output from the graph represented by the node named `add_tensor_2`
    buffers_to_mutate={'add_tensor_2': 'L__self___my_buffer2'},

    # Backward graph not captured
    backward_signature=None,

    # Work in progress feature, please ignore now.
    assertion_dep_token=None
)

class torch.export.ArgumentKind(value) [source]

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

class torch.export.ArgumentSpec(kind: torch.export.ArgumentKind, value: Any) [source]

class torch.export.ModuleCallSignature(inputs: List[torch.export.ArgumentSpec], outputs: List[torch.export.ArgumentSpec], in_spec: torch.utils._pytree.TreeSpec, out_spec: torch.utils._pytree.TreeSpec) [source]

class torch.export.ModuleCallEntry(fqn: str, signature: Union[torch.export.ModuleCallSignature, NoneType] = None) [source]