Введение в TorchScript

глубокое обучение

Это руководство представляет собой введение в TorchScript, модель PyTorch (nn.ModuleПодкласс ), который может работать в высокопроизводительной среде, такой как C++.

В этом уроке мы рассмотрим:

  1. Основы разработки моделей в PyTorch, в том числе:
    • модуль
    • определить прямую функцию
    • Группировать модули в иерархию модулей
  2. Конкретный способ преобразования модулей PyTorch в TorchScript (наша высокопроизводительная среда выполнения развертывания)
    • отслеживать существующие модули
    • Компилировать модули напрямую с помощью скриптов
    • Как совместить два метода
    • Сохранение и загрузка модулей TorchScript

Мы надеемся, что после завершения этого руководства вы перейдете к следующему руководству, в котором будет рассмотрен пример фактического вызова модели TorchScript из C++.

import torch  # 这是同时使用PyTorch和TorchScript所需的全部导入!
print(torch.__version__)

  • выходной результат
1.3.0

1. Основа модели PyTorch

Начнем с определения простого модуля. Модуль — это основная единица композиции в PyTorch. Это содержит:

  • Конструктор, подготавливающий модуль к вызову
  • Набор параметров и подмодулей. Они инициализируются конструктором и могут использоваться модулями во время вызова.
  • функция вперед. Это код, который запускается при вызове модуля.
    Давайте рассмотрим небольшой пример:
class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()

    def forward(self, x, h):
        new_h = torch.tanh(x + h)
        return new_h, new_h

my_cell = MyCell()
x = torch.rand(3, 4)
h = torch.rand(3, 4)
print(my_cell(x, h))

  • выходной результат
(tensor([[0.5139, 0.6451, 0.3697, 0.7738],
        [0.7936, 0.5864, 0.8063, 0.9324],
        [0.6479, 0.8408, 0.8062, 0.7263]]), tensor([[0.5139, 0.6451, 0.3697, 0.7738],
        [0.7936, 0.5864, 0.8063, 0.9324],
        [0.6479, 0.8408, 0.8062, 0.7263]]))

Таким образом, мы имеем:

  1. создал подклассtorch.nn.Moduleтип.
  2. Определите конструктор. Конструктор мало что делает, он просто вызывает конструктор super.
  3. Определена прямая функция, которая принимает два входа и возвращает два выхода. Фактическое содержание прямой функции не очень важно, но это подделка.блок РНН
    - То есть функция применяется к циклу.

Мы создаем экземпляр модуля и делаемxиy, это просто матрицы 3x4 случайных значений. Затем мы используемmy_cell(x,h)позвонить на сотовый. Это снова вызывает нашу функцию пересылки.

Давайте сделаем что-нибудь более интересное:

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

  • выходной результат
MyCell(
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[ 0.3941,  0.4160, -0.1086,  0.8432],
        [ 0.5604,  0.4003,  0.5009,  0.6842],
        [ 0.7084,  0.7147,  0.1818,  0.8296]], grad_fn=<TanhBackward>), tensor([[ 0.3941,  0.4160, -0.1086,  0.8432],
        [ 0.5604,  0.4003,  0.5009,  0.6842],
        [ 0.7084,  0.7147,  0.1818,  0.8296]], grad_fn=<TanhBackward>))

Мы переопределили модульMyCell, но на этот раз мы добавилиself.linearатрибут и вызывается в прямой функцииself.linear.

Что, черт возьми, здесь происходит?torch.nn.Linear— это модуль стандартной библиотеки PyTorch. какMyCellОпять же, его можно вызвать, используя синтаксис вызова. Мы строим иерархию модулей.

Печать на модуле обеспечивает визуальное представление иерархии подклассов этого модуля. В нашем примере мы можем видеть наш линейный подкласс и его параметры.

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

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

Теперь давайте проверим его гибкость:

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.dg = MyDecisionGate()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell()
print(my_cell)
print(my_cell(x, h))

  • выходной результат
MyCell(
  (dg): MyDecisionGate()
  (linear): Linear(in_features=4, out_features=4, bias=True)
)
(tensor([[0.0850, 0.2812, 0.5188, 0.8523],
        [0.1233, 0.3948, 0.6615, 0.7466],
        [0.7072, 0.6103, 0.6953, 0.7047]], grad_fn=<TanhBackward>), tensor([[0.0850, 0.2812, 0.5188, 0.8523],
        [0.1233, 0.3948, 0.6615, 0.7466],
        [0.7072, 0.6103, 0.6953, 0.7047]], grad_fn=<TanhBackward>))

мы переопределяемMyCellкласс, но здесь мы определяемMyDecisionGate. Этот модуль использует поток управления. Поток управления включает в себя такие вещи, как циклы и операторы if.

Учитывая полное представление программы, многие фреймворки используют подход вычисления символьного деривации. Однако в PyTorch мы используем полосы градиента. Мы записываем действия по мере их возникновения и воспроизводим их при вычислении производных. Таким образом, платформе не нужно явно определять производные классы для всех конструкций языка.

2. Основы TorchScript

Теперь давайте возьмем работающий пример и посмотрим, как можно применить TorchScript.

Короче говоря, несмотря на то, что PyTorch является гибким и динамичным, TorchScript предоставляет инструменты для захвата определений моделей. Начнем с так называемогоотслеживать.

2.1 Модуль трассировки

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

  • выходной результат
TracedModule[MyCell](
  original_name=MyCell
  (linear): TracedModule[Linear](original_name=Linear)
)

Мы сделали шаг назад и выбралиMyCellВторой вариант класса. Как и раньше, мы создаем его экземпляр, но на этот раз мы вызываемtorch.jit.trace,существуетModule(модуль) передается пример и входные данные, которые сеть может увидеть в примере.

Что именно это делает? Он вызвал модуль, записал, что произошло во время работы модуля, и создалtorch.jit.ScriptModuleэкземпляр (TracedModuleэто его экземпляр)

TorchScript записывает свое определение в промежуточном представлении (или IR), обычно называемом графом в глубоком обучении. Мы можем проверить с.graphГрафик свойств:

print(traced_cell.graph)

  • выходной результат
graph(%self : ClassType<MyCell>,
      %input : Float(3, 4),
      %h : Float(3, 4)):
  %1 : ClassType<Linear> = prim::GetAttr[name="linear"](%self)
  %weight : Tensor = prim::GetAttr[name="weight"](%1)
  %bias : Tensor = prim::GetAttr[name="bias"](%1)
  %6 : Float(4, 4) = aten::t(%weight), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0
  %7 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0
  %8 : int = prim::Constant[value=1](), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0
  %9 : Float(3, 4) = aten::addmm(%bias, %input, %6, %7, %8), scope: MyCell/Linear[linear] # /opt/conda/lib/python3.6/site-packages/torch/nn/functional.py:1370:0
  %10 : int = prim::Constant[value=1](), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %11 : Float(3, 4) = aten::add(%9, %h, %10), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %12 : Float(3, 4) = aten::tanh(%11), scope: MyCell # /var/lib/jenkins/workspace/beginner_source/Intro_to_TorchScript_tutorial.py:188:0
  %13 : (Float(3, 4), Float(3, 4)) = prim::TupleConstruct(%12, %12)
  return (%13)

Однако это очень низкоуровневое представление, и большая часть информации, содержащейся в графе, бесполезна для конечного пользователя. Вместо этого мы можем использовать.codeатрибут, чтобы дать синтаксическую интерпретацию кода Python:

print(traced_cell.code)

  • выходной результат
import __torch__
import __torch__.torch.nn.modules.linear
def forward(self,
    input: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = self.linear
  weight = _0.weight
  bias = _0.bias
  _1 = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
  _2 = torch.tanh(torch.add(_1, h, alpha=1))
  return (_2, _2)

Так зачем мы все это делаем? Причин несколько:

  1. Код TorchScript можно вызывать в собственном интерпретаторе, который по сути является ограниченным интерпретатором Python. Интерпретатор не заблокирован глобальным интерпретатором, поэтому его можно
    Многие запросы обрабатываются одновременно на одном экземпляре.
  2. Этот формат позволяет нам сохранить всю модель на диск и загрузить ее в другую среду, например, на сервер, написанный на языке, отличном от Python.
  3. TorchScript предоставляет нам представление, в котором мы можем выполнять оптимизацию компилятора нашего кода, чтобы обеспечить более эффективное выполнение.
  4. TorchScript позволяет нам взаимодействовать со многими средами выполнения бэкэнда/устройства, которые требуют более широкого представления программы, чем один оператор.

Мы видим, что вызовtraced_cellдает тот же результат, что и модуль Python:

print(my_cell(x, h))
print(traced_cell(x, h))

  • выходной результат
(tensor([[-0.3983,  0.5954,  0.2587, -0.3748],
        [-0.5033,  0.4471,  0.8264,  0.2135],
        [ 0.3430,  0.5561,  0.6794, -0.2273]], grad_fn=<TanhBackward>), tensor([[-0.3983,  0.5954,  0.2587, -0.3748],
        [-0.5033,  0.4471,  0.8264,  0.2135],
        [ 0.3430,  0.5561,  0.6794, -0.2273]], grad_fn=<TanhBackward>))
(tensor([[-0.3983,  0.5954,  0.2587, -0.3748],
        [-0.5033,  0.4471,  0.8264,  0.2135],
        [ 0.3430,  0.5561,  0.6794, -0.2273]],
       grad_fn=<DifferentiableGraphBackward>), tensor([[-0.3983,  0.5954,  0.2587, -0.3748],
        [-0.5033,  0.4471,  0.8264,  0.2135],
        [ 0.3430,  0.5561,  0.6794, -0.2273]],
       grad_fn=<DifferentiableGraphBackward>))

3. Используйте скрипт для преобразования модуля

Есть причина, по которой мы используем вторую версию модуля вместо подмодулей с большим потоком управления. Теперь проверим:

class MyDecisionGate(torch.nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x
        else:
            return -x

class MyCell(torch.nn.Module):
    def __init__(self, dg):
        super(MyCell, self).__init__()
        self.dg = dg
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.dg(self.linear(x)) + h)
        return new_h, new_h

my_cell = MyCell(MyDecisionGate())
traced_cell = torch.jit.trace(my_cell, (x, h))
print(traced_cell.code)

  • выходной результат
import __torch__.___torch_mangle_0
import __torch__
import __torch__.torch.nn.modules.linear.___torch_mangle_1
def forward(self,
    input: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = self.linear
  weight = _0.weight
  bias = _0.bias
  x = torch.addmm(bias, input, torch.t(weight), beta=1, alpha=1)
  _1 = torch.tanh(torch.add(x, h, alpha=1))
  return (_1, _1)

Проверять.codeвывод, мы можем найти, где не найденоif-elseФилиал! Зачем? Трассировка делает именно то, о чем мы говорим: запускает код, регистрирует происходящее и создает систему, которая может это делать.ScriptModule. К сожалению, такие вещи, как поток управления, уничтожены.

Как мы можем точно представить этот модуль в TorchScript? мы предоставляемкомпилятор сценариев, который напрямую анализирует ваш исходный код Python, чтобы преобразовать его в TorchScript. Воспользуемся компилятором скриптов для преобразованияMyDecisionGate:

scripted_gate = torch.jit.script(MyDecisionGate())

my_cell = MyCell(scripted_gate)
traced_cell = torch.jit.script(my_cell)
print(traced_cell.code)

  • выходной результат
import __torch__.___torch_mangle_3
import __torch__.___torch_mangle_2
import __torch__.torch.nn.modules.linear.___torch_mangle_4
def forward(self,
    x: Tensor,
    h: Tensor) -> Tuple[Tensor, Tensor]:
  _0 = self.linear
  _1 = _0.weight
  _2 = _0.bias
  if torch.eq(torch.dim(x), 2):
    _3 = torch.__isnot__(_2, None)
  else:
    _3 = False
  if _3:
    bias = ops.prim.unchecked_unwrap_optional(_2)
    ret = torch.addmm(bias, x, torch.t(_1), beta=1, alpha=1)
  else:
    output = torch.matmul(x, torch.t(_1))
    if torch.__isnot__(_2, None):
      bias0 = ops.prim.unchecked_unwrap_optional(_2)
      output0 = torch.add_(output, bias0, alpha=1)
    else:
      output0 = output
    ret = output0
  _4 = torch.gt(torch.sum(ret, dtype=None), 0)
  if bool(_4):
    _5 = ret
  else:
    _5 = torch.neg(ret)
  new_h = torch.tanh(torch.add(_5, h, alpha=1))
  return (new_h, new_h)

Теперь мы точно зафиксировали поведение нашей программы в TorchScript. Теперь попробуем запустить программу:

# New inputs
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell(x, h)

3.1 Сочетание сценариев и трассировки

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

  • Пример первого случая:
class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super(MyRNNLoop, self).__init__()
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))

    def forward(self, xs):
        h, y = torch.zeros(3, 4), torch.zeros(3, 4)
        for i in range(xs.size(0)):
            y, h = self.cell(xs[i], h)
        return y, h

rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)

  • выходной результат
import __torch__
import __torch__.___torch_mangle_5
import __torch__.___torch_mangle_2
import __torch__.torch.nn.modules.linear.___torch_mangle_6
def forward(self,
    xs: Tensor) -> Tuple[Tensor, Tensor]:
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  y0 = y
  h0 = h
  for i in range(torch.size(xs, 0)):
    _0 = self.cell
    _1 = torch.select(xs, 0, i)
    _2 = _0.linear
    weight = _2.weight
    bias = _2.bias
    _3 = torch.addmm(bias, _1, torch.t(weight), beta=1, alpha=1)
    _4 = torch.gt(torch.sum(_3, dtype=None), 0)
    if bool(_4):
      _5 = _3
    else:
      _5 = torch.neg(_3)
    _6 = torch.tanh(torch.add(_5, h0, alpha=1))
    y0, h0 = _6, _6
  return (y0, h0)

  • Пример второго случая:
class WrapRNN(torch.nn.Module):
    def __init__(self):
        super(WrapRNN, self).__init__()
        self.loop = torch.jit.script(MyRNNLoop())

    def forward(self, xs):
        y, h = self.loop(xs)
        return torch.relu(y)

traced = torch.jit.trace(WrapRNN(), (torch.rand(10, 3, 4)))
print(traced.code)

  • выходной результат
import __torch__
import __torch__.___torch_mangle_9
import __torch__.___torch_mangle_7
import __torch__.___torch_mangle_2
import __torch__.torch.nn.modules.linear.___torch_mangle_8
def forward(self,
    argument_1: Tensor) -> Tensor:
  _0 = self.loop
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  h0 = h
  for i in range(torch.size(argument_1, 0)):
    _1 = _0.cell
    _2 = torch.select(argument_1, 0, i)
    _3 = _1.linear
    weight = _3.weight
    bias = _3.bias
    _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
    _5 = torch.gt(torch.sum(_4, dtype=None), 0)
    if bool(_5):
      _6 = _4
    else:
      _6 = torch.neg(_4)
    h0 = torch.tanh(torch.add(_6, h0, alpha=1))
  return torch.relu(h0)

Таким образом, сценарии и трассировки можно использовать вместе, когда того требует ситуация.

4. Сохраните и загрузите модель

Мы предоставляем API для сохранения или загрузки модулей TorchScript на диск в формате архива. Этот формат включает в себя код, параметры, свойства и отладочную информацию, что означает, что архив является независимым представлением модели, которое может быть загружено в совершенно независимом процессе. Давайте сохраним и загрузим завернутыйRNNМодуль:

traced.save('wrapped_rnn.zip')

loaded = torch.jit.load('wrapped_rnn.zip')

print(loaded)
print(loaded.code)

  • выходной результат
ScriptModule(
  original_name=WrapRNN
  (loop): ScriptModule(
    original_name=MyRNNLoop
    (cell): ScriptModule(
      original_name=MyCell
      (dg): ScriptModule(original_name=MyDecisionGate)
      (linear): ScriptModule(original_name=Linear)
    )
  )
)
import __torch__
import __torch__.___torch_mangle_9
import __torch__.___torch_mangle_7
import __torch__.___torch_mangle_2
import __torch__.torch.nn.modules.linear.___torch_mangle_8
def forward(self,
    argument_1: Tensor) -> Tensor:
  _0 = self.loop
  h = torch.zeros([3, 4], dtype=None, layout=None, device=None, pin_memory=None)
  h0 = h
  for i in range(torch.size(argument_1, 0)):
    _1 = _0.cell
    _2 = torch.select(argument_1, 0, i)
    _3 = _1.linear
    weight = _3.weight
    bias = _3.bias
    _4 = torch.addmm(bias, _2, torch.t(weight), beta=1, alpha=1)
    _5 = torch.gt(torch.sum(_4, dtype=None), 0)
    if bool(_5):
      _6 = _4
    else:
      _6 = torch.neg(_4)
    h0 = torch.tanh(torch.add(_6, h0, alpha=1))
  return torch.relu(h0)

Как видите, сериализация сохраняет иерархию модулей и код, над которым мы работали. Например, модели также можно загружать в C++ для выполнения независимо от Python.

дальнейшее чтение

Мы закончили с учебником! Чтобы получить более подробную демонстрацию, ознакомьтесь с демонстрацией NeurIPS для преобразования моделей машинного перевода с помощью TorchScript: https://colab.research.google.com/drive/1HiICg6jRkBnr5hvK2-VnMi88Vi9pUzEJ.

Общее время работы скрипта: (0 минут 0,247 секунды)

Сводная станция блога о технологиях искусственного интеллекта Panchuang: http://docs.panchuang.net/PyTorch, официальная учебная станция на китайском языке: http://pytorch.panchuang.net/OpenCV, официальный китайский документ: http://woshicver.com/