Модель программирования torch.export

Строгая и нестрогая трассировка

torch.export.export() предоставляет два режима трассировки.

В нестрогом режиме мы трассируем выполнение программы с помощью обычного интерпретатора Python. Ваш код выполняется в точности так же, как в eager-режиме; единственное отличие заключается в том, что все тензоры заменяются на фиктивные тензоры, которые имеют формы и другие виды метаданных, но не содержат данных, и оборачиваются в объекты Proxy, записывающие все выполняемые над ними операции в граф. Мы также захватываем условия для форм тензоров, обеспечивающие корректность сгенерированного кода.

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

Хотя в настоящее время режим трассировки по умолчанию — строгий, мы настоятельно рекомендуем использовать нестрогий режим, который вскоре станет режимом по умолчанию. Для большинства моделей условий для форм тензоров достаточно, чтобы обеспечить корректность, а дополнительные гарантии безопасности на уровне Python не имеют значения; в то же время возможность столкнуться с неподдерживаемыми возможностями Python в TorchDynamo создаёт ненужный риск.

В остальной части этого документа предполагается, что трассировка выполняется в нестрогом режиме; в частности, предполагается, что поддерживаются все возможности Python.

Статические значения

Статическое значение — это значение, которое зафиксировано во время экспорта и не может меняться при выполнении экспортированной программы. При обнаружении такого значения во время трассировки мы считаем его константой и встраиваем в граф.

Если выполняется операция (например, x + y) и все её входные данные статические, результат операции напрямую встраивается в граф, а сама операция в нём не отображается (то есть происходит «свёртка констант»).

Если значение встроено в граф, мы говорим, что граф специализирован под это значение. Например:

import torch

class MyMod(torch.nn.Module):
    def forward(self, x, y):
        z = y + 7
        return x + z

m = torch.export.export(MyMod(), (torch.randn(1), 3))
print(m.graph_module.code)

"""
def forward(self, arg0_1, arg1_1):
    add = torch.ops.aten.add.Tensor(arg0_1, 10);  arg0_1 = None
    return (add,)

"""

Здесь в качестве трассируемого значения для y передаётся 3; оно рассматривается как статическое значение и добавляется к 7, встраивая статическое значение 10 в граф.

Динамические значения

Динамическое значение — это значение, которое может меняться при каждом запуске. Оно ведёт себя как «обычный» аргумент функции: можно передавать разные входные данные и ожидать, что функция будет работать корректно.

Какие значения являются статическими, а какие — динамическими?

Статическим или динамическим значение считается в зависимости от его типа:

• Для тензоров:

  • Данные тензора считаются динамическими.

  • Формы тензоров система может считать статическими или динамическими.

    • По умолчанию формы всех входных тензоров считаются статическими. Пользователь может переопределить это поведение для любого входного тензора, задав для него динамическую форму.
    • Формы тензоров, входящих в состояние модуля, то есть параметров и буферов, всегда статические.
  • Другие виды метаданных тензоров (например, device, dtype) статические.

• Примитивы Python (int, float, bool, str, None) являются статическими.

  • Для некоторых примитивных типов существуют динамические варианты (SymInt, SymFloat, SymBool). Обычно пользователям не нужно с ними работать.
  • Пользователи могут указать, что целочисленные входные данные являются динамическими, задав для них динамическую форму.

• Для стандартных контейнеров Python (list, tuple, dict, namedtuple):

  • Структура (то есть длина для значений list и tuple, а также последовательность ключей для значений dict и namedtuple) является статической.
  • К содержащимся в контейнерах элементам рекурсивно применяются эти правила (по сути, схема PyTree), причём листьями являются тензоры или примитивные типы.
• Другие классы (в том числе классы данных) можно зарегистрировать как PyTree (см. ниже); к ним применяются те же правила, что и к стандартным контейнерам.

Пользовательские типы входных данных (PyTree)

Кроме того, вы можете определить собственный (пользовательский) класс и использовать его как тип входных данных, но для этого потребуется зарегистрировать такой класс как PyTree.

Ниже приведён пример использования утилиты для регистрации класса данных, используемого в качестве типа входных данных.

@dataclass
class Input:
    f: torch.Tensor
    p: torch.Tensor

import torch.utils._pytree as pytree
pytree.register_dataclass(Input)

class M(torch.nn.Module):
    def forward(self, x: Input):
        return x.f + 1

torch.export.export(M(), (Input(f=torch.ones(10, 4), p=torch.zeros(10, 4)),))

Необязательные типы входных данных

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

class M(torch.nn.Module):
    def forward(self, x, y=None):
        if y is not None:
            return y * x
        return x + x

# Optional input is passed in
ep = torch.export.export(M(), (torch.randn(3, 3), torch.randn(3, 3)))
print(ep)
"""
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 3]", y: "f32[3, 3]"):
            # File: /data/users/angelayi/pytorch/moo.py:15 in forward, code: return y * x
            mul: "f32[3, 3]" = torch.ops.aten.mul.Tensor(y, x);  y = x = None
            return (mul,)
"""

# Optional input is not passed in
ep = torch.export.export(M(), (torch.randn(3, 3),))
print(ep)
"""
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, x: "f32[3, 3]", y):
            # File: /data/users/angelayi/pytorch/moo.py:16 in forward, code: return x + x
            add: "f32[3, 3]" = torch.ops.aten.add.Tensor(x, x);  x = None
            return (add,)
"""

Статический поток управления

Поток управления Python на основе статических значений поддерживается автоматически. (Напомним, что к статическим значениям относятся статические формы, поэтому поток управления на основе статических форм также относится к этому случаю.)

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

В случае оператора if мы продолжим трассировку ветви, выбранной во время экспорта. В случае оператора for или while мы продолжим трассировку, разворачивая цикл.

Динамический поток управления: зависимый от формы и зависимый от данных

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

class M_old(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x.sin()
        else:
            return x.cos()

class M_new(torch.nn.Module):
    def forward(self, x):
        return torch.cond(
            pred=x.sum() > 0,
            true_fn=lambda x: x.sin(),
            false_fn=lambda x: x.cos(),
            operands=(x,),
        )

class M_old(torch.nn.Module):
    def forward(self, x):
        nz = x.nonzero()
        if nz.shape[0] > 0:
            return x.sin()
        else:
            return x.cos()

class M_new(torch.nn.Module):
    def forward(self, x):
        nz = x.nonzero()
        torch._check(nz.shape[0] > 0)
        if nz.shape[0] > 0:
            return x.sin()
        else:
            return x.cos()

Фиктивные реализации операторов PyTorch

Напомним, что во время трассировки программа выполняется с использованием фиктивных тензоров, не содержащих данных. Как правило, фактические реализации операторов PyTorch нельзя вызывать с фиктивными тензорами. Поэтому для каждого оператора необходима дополнительная фиктивная (или «мета-») реализация, принимающая фиктивные тензоры и возвращающая фиктивные тензоры и соответствующая поведению фактической реализации по формам и другим видам метаданных, содержащихся в фиктивных тензорах.

Например, обратите внимание, как фиктивная реализация torch.index_select() вычисляет форму результата по форме входных данных (игнорируя входные данные и возвращая пустые выходные данные).

def meta_index_select(self, dim, index):
    result_size = list(self.size())
    if self.dim() > 0:
        result_size[dim] = index.numel()
    return self.new_empty(result_size)

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

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

• Если условие включает только динамические формы с опорой, мы можем использовать их конкретные значения, чтобы определить, равно ли условие True или False. Затем мы можем добавить в среду форм защитное условие, указывающее, что соответствующее символьное булево выражение равно True или False, и продолжить трассировку.
• В противном случае условие включает динамические формы без опоры. Как правило, без дополнительной информации вычислить такое условие невозможно; поэтому мы не можем продолжить трассировку и должны вызвать ошибку во время экспорта. Для продолжения трассировки пользователю следует использовать явный оператор PyTorch. Эта информация добавляется в среду форм как защитное условие и также может помочь вычислить другие встречающиеся далее условия, определить, равны ли они True или False.

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

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

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

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

Ниже приведён пример пользовательского оператора sin, использующего NumPy, и его зарегистрированной (тривиальной) фиктивной реализации.

@torch.library.custom_op("mylib::sin", mutates_args=())
def sin(x: Tensor) -> Tensor:
    x_np = x.numpy()
    y_np = np.sin(x_np)
    return torch.from_numpy(y_np)

@torch.library.register_fake("mylib::sin")
def _(x: Tensor) -> Tensor:
    return torch.empty_like(x)

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

...

@torch.library.register_fake("mylib::custom_nonzero")
def _(x):
    nnz = torch.library.get_ctx().new_dynamic_size()
    shape = [nnz, x.dim()]
    return x.new_empty(shape, dtype=torch.int64)

Правила доступа

Все состояния модуля должны быть инициализированы. Обращение к неинициализированному состоянию модуля приводит к ошибке во время экспорта.

Чтение состояния модуля всегда разрешено.

Обновление состояния модуля возможно, но должно соответствовать следующим правилам:

• Статический обычный атрибут (например, примитивного типа) можно обновлять. Чтение и обновление можно свободно чередовать; как и ожидается, при чтении всегда будут видны значения последних обновлений. Поскольку эти атрибуты статические, их значения также встраиваются, поэтому сгенерированный код не будет содержать инструкций для фактического «получения» или «установки» таких атрибутов.
• Динамический обычный атрибут (например, тензорного типа) обновлять нельзя. Для этого его необходимо зарегистрировать как буфер во время инициализации модуля.
• Буфер можно обновлять — как на месте (например, self.buffer[:] = ...), так и не на месте (например, self.buffer = ...).
• Параметр обновлять нельзя. Обычно параметры обновляются только во время обучения, а не во время инференса. Чтобы избежать обновления параметров во время экспорта, рекомендуется выполнять экспорт с помощью torch.no_grad().

Эффекты функционализации

Любое динамическое состояние модуля, к которому выполняется чтение и/или которое обновляется, «выносится» в качестве входных данных и/или выходных данных сгенерированного кода соответственно.

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