Это руководство представляет собой введение в TorchScript, модель PyTorch (nn.Module
Подкласс ), который может работать в высокопроизводительной среде, такой как C++.
В этом уроке мы рассмотрим:
- Основы разработки моделей в PyTorch, в том числе:
- модуль
- определить прямую функцию
- Группировать модули в иерархию модулей
- Конкретный способ преобразования модулей 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]]))
Таким образом, мы имеем:
- создал подкласс
torch.nn.Module
тип. - Определите конструктор. Конструктор мало что делает, он просто вызывает конструктор super.
- Определена прямая функция, которая принимает два входа и возвращает два выхода. Фактическое содержание прямой функции не очень важно, но это подделка.блок РНН
- То есть функция применяется к циклу.
Мы создаем экземпляр модуля и делаем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)
Так зачем мы все это делаем? Причин несколько:
- Код TorchScript можно вызывать в собственном интерпретаторе, который по сути является ограниченным интерпретатором Python. Интерпретатор не заблокирован глобальным интерпретатором, поэтому его можно
Многие запросы обрабатываются одновременно на одном экземпляре. - Этот формат позволяет нам сохранить всю модель на диск и загрузить ее в другую среду, например, на сервер, написанный на языке, отличном от Python.
- TorchScript предоставляет нам представление, в котором мы можем выполнять оптимизацию компилятора нашего кода, чтобы обеспечить более эффективное выполнение.
- 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/