TorchScript

Исключительные случаи трассировки

Существуют некоторые исключительные случаи, когда трассировка данной Python-функции/модуля не будет отражать базовый код. Эти случаи могут включать:

• Трассировку управляющей последовательности, зависящей от входных данных (например, формы тензоров)
• Трассировку операций на месте для представлений тензоров (например, индексирования в левой части присваивания)

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

Автоматическая проверка трассировки

Один из способов автоматически обнаруживать многие ошибки в трассировках — использовать check_inputs в API torch.jit.trace(). check_inputs принимает список кортежей входных данных, которые будут использованы для повторной трассировки вычислений и проверки результатов. Например:

def loop_in_traced_fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

traced = torch.jit.trace(loop_in_traced_fn, inputs, check_inputs=check_inputs)

Что даёт нам следующую диагностическую информацию:

ERROR: Graphs differed across invocations!
Graph diff:

            graph(%x : Tensor) {
            %1 : int = prim::Constant[value=0]()
            %2 : int = prim::Constant[value=0]()
            %result.1 : Tensor = aten::select(%x, %1, %2)
            %4 : int = prim::Constant[value=0]()
            %5 : int = prim::Constant[value=0]()
            %6 : Tensor = aten::select(%x, %4, %5)
            %result.2 : Tensor = aten::mul(%result.1, %6)
            %8 : int = prim::Constant[value=0]()
            %9 : int = prim::Constant[value=1]()
            %10 : Tensor = aten::select(%x, %8, %9)
        -   %result : Tensor = aten::mul(%result.2, %10)
        +   %result.3 : Tensor = aten::mul(%result.2, %10)
        ?          ++
            %12 : int = prim::Constant[value=0]()
            %13 : int = prim::Constant[value=2]()
            %14 : Tensor = aten::select(%x, %12, %13)
        +   %result : Tensor = aten::mul(%result.3, %14)
        +   %16 : int = prim::Constant[value=0]()
        +   %17 : int = prim::Constant[value=3]()
        +   %18 : Tensor = aten::select(%x, %16, %17)
        -   %15 : Tensor = aten::mul(%result, %14)
        ?     ^                                 ^
        +   %19 : Tensor = aten::mul(%result, %18)
        ?     ^                                 ^
        -   return (%15);
        ?             ^
        +   return (%19);
        ?             ^
            }

Это сообщение указывает нам на то, что вычисление отличается между первой трассировкой и трассировкой с check_inputs. Действительно, цикл внутри тела loop_in_traced_fn зависит от формы входного тензора x, и поэтому при другой трассировке с другой формой граф отличается.

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

def fn(x):
    result = x[0]
    for i in range(x.size(0)):
        result = result * x[i]
    return result

inputs = (torch.rand(3, 4, 5),)
check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

scripted_fn = torch.jit.script(fn)
print(scripted_fn.graph)
#print(str(scripted_fn.graph).strip())

for input_tuple in [inputs] + check_inputs:
    torch.testing.assert_close(fn(*input_tuple), scripted_fn(*input_tuple))

Что даёт:

graph(%x : Tensor) {
    %5 : bool = prim::Constant[value=1]()
    %1 : int = prim::Constant[value=0]()
    %result.1 : Tensor = aten::select(%x, %1, %1)
    %4 : int = aten::size(%x, %1)
    %result : Tensor = prim::Loop(%4, %5, %result.1)
    block0(%i : int, %7 : Tensor) {
        %10 : Tensor = aten::select(%x, %1, %i)
        %result.2 : Tensor = aten::mul(%7, %10)
        -> (%5, %result.2)
    }
    return (%result);
}

Предупреждения трассера

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

def fill_row_zero(x):
    x[0] = torch.rand(*x.shape[1:2])
    return x

traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
print(traced.graph)

Генерирует несколько предупреждений и граф, просто возвращающий входной тензор:

fill_row_zero.py:4: TracerWarning: There are 2 live references to the data region being modified when tracing in-place operator copy_ (possibly due to an assignment). This might cause the trace to be incorrect, because all other views that also reference this data will not reflect this change in the trace! On the other hand, if all other views use the same memory chunk, but are disjoint (e.g. are outputs of torch.split), this might still be safe.
    x[0] = torch.rand(*x.shape[1:2])
fill_row_zero.py:6: TracerWarning: Output nr 1. of the traced function does not match the corresponding output of the Python function. Detailed error:
Not within tolerance rtol=1e-05 atol=1e-05 at input[0, 1] (0.09115803241729736 vs. 0.6782537698745728) and 3 other locations (33.00%)
    traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
graph(%0 : Float(3, 4)) {
    return (%0);
}

Это можно исправить, изменив код так, чтобы не использовать обновление на месте, а вместо этого собирать результирующий тензор вне места с torch.cat:

def fill_row_zero(x):
    x = torch.cat((torch.rand(1, *x.shape[1:2]), x[1:2]), dim=0)
    return x

traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
print(traced.graph)

Модули

При передаче в функцию torch.jit.script, данные torch.nn.Module копируются в ScriptModule, и компилятор TorchScript компилирует модуль. Модуль по умолчанию компилирует свой forward. Методы, вызываемые из forward, компилируются лениво в том порядке, в котором они используются в forward, а также любые методы @torch.jit.export.

torch.jit.export(fn) [source]

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

forward неявно предполагается, что является точкой входа, поэтому для него не нужен этот декоратор. Функции и методы, вызываемые из forward , компилируются по мере их обнаружения компилятором, поэтому для них также не нужен этот декоратор.

Пример (использование @torch.jit.export для метода):

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def implicitly_compiled_method(self, x):
        return x + 99

    # `forward` is implicitly decorated with `@torch.jit.export`,
    # so adding it here would have no effect
    def forward(self, x):
        return x + 10

    @torch.jit.export
    def another_forward(self, x):
        # When the compiler sees this call, it will compile
        # `implicitly_compiled_method`
        return self.implicitly_compiled_method(x)

    def unused_method(self, x):
        return x - 20

# `m` will contain compiled methods:
#     `forward`
#     `another_forward`
#     `implicitly_compiled_method`
# `unused_method` will not be compiled since it was not called from
# any compiled methods and wasn't decorated with `@torch.jit.export`
m = torch.jit.script(MyModule())

Функции

Функции практически не меняются, их можно декорировать @torch.jit.ignore или torch.jit.unused при необходимости.

# Same behavior as pre-PyTorch 1.2
@torch.jit.script
def some_fn():
    return 2

# Marks a function as ignored, if nothing
# ever calls it then this has no effect
@torch.jit.ignore
def some_fn2():
    return 2

# As with ignore, if nothing calls it then it has no effect.
# If it is called in script it is replaced with an exception.
@torch.jit.unused
def some_fn3():
  import pdb; pdb.set_trace()
  return 4

# Doesn't do anything, this function is already
# the main entry point
@torch.jit.export
def some_fn4():
    return 2

Классы TorchScript

Все в пользовательском классе TorchScript экспортируется по умолчанию; функции могут быть декорированы @torch.jit.ignore при необходимости.

Атрибуты

Компилятору TorchScript необходимо знать типы module attributes. Большинство типов могут быть выведены из значения члена. Пустые списки и словари не могут иметь тип выведен и должны иметь тип, указанный аннотациями классов в стиле PEP 526. Если тип не может быть выведен и не указан явно, он не будет добавлен в качестве атрибута в результирующий ScriptModule

Старый API:

from typing import Dict
import torch

class MyModule(torch.jit.ScriptModule):
    def __init__(self):
        super(MyModule, self).__init__()
        self.my_dict = torch.jit.Attribute({}, Dict[str, int])
        self.my_int = torch.jit.Attribute(20, int)

m = MyModule()

Новый API:

from typing import Dict

class MyModule(torch.nn.Module):
    my_dict: Dict[str, int]

    def __init__(self):
        super(MyModule, self).__init__()
        # This type cannot be inferred and must be specified
        self.my_dict = {}

        # The attribute type here is inferred to be `int`
        self.my_int = 20

    def forward(self):
        pass

m = torch.jit.script(MyModule())

Константы

Конструктор типа Final может быть использован для маркировки членов как constant. Если члены не помечены как константы, они будут скопированы в результирующий ScriptModule как атрибут. Использование Final открывает возможности для оптимизации, если значение известно как фиксированное, и обеспечивает дополнительную безопасность типа.

Старый API:

class MyModule(torch.jit.ScriptModule):
    __constants__ = ['my_constant']

    def __init__(self):
        super(MyModule, self).__init__()
        self.my_constant = 2

    def forward(self):
        pass
m = MyModule()

Новый API:

try:
    from typing_extensions import Final
except:
    # If you don't have `typing_extensions` installed, you can use a
    # polyfill from `torch.jit`.
    from torch.jit import Final

class MyModule(torch.nn.Module):

    my_constant: Final[int]

    def __init__(self):
        super(MyModule, self).__init__()
        self.my_constant = 2

    def forward(self):
        pass

m = torch.jit.script(MyModule())

Переменные

Контейнеры предполагаются иметь тип Tensor и не быть необязательными (см. Default Types для получения дополнительной информации). Ранее использовалось torch.jit.annotate для указания компилятору TorchScript, каким должен быть тип. Теперь поддерживаются указания типа в стиле Python 3.

import torch
from typing import Dict, Optional

@torch.jit.script
def make_dict(flag: bool):
    x: Dict[str, int] = {}
    x['hi'] = 2
    b: Optional[int] = None
    if flag:
        b = 2
    return x, b