Загрузка моделей TorchScript в C++

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

Это руководство было обновлено для работы с PyTorch 1.2.

Как следует из названия, основным интерфейсом PyTorch является язык программирования Python. Хотя Python подходит и является предпочтительным языком для многих сценариев, требующих динамизма и простоты итерации, также существует множество ситуаций, когда эти свойства Python просто неблагоприятны. Одна среда, где обычно применяется последнее, — это среда, в которой требуется производство — низкая задержка и строгое развертывание. Для производственных сценариев часто предпочтительным языком является C++, даже если он привязан только к другому языку, такому как Java, Rust или Go. В следующих абзацах будет описан путь, предоставляемый PyTorch, от существующей модели Python к сериализованному представлению, которое можно загрузить и выполнить полностью из C++, не полагаясь на Python.

Шаг 1. Преобразование модели PyTorch в скрипт Torch

Переход модели PyTorch от Python к C++ инициируется Torch Script, представлением модели PyTorch, которое может быть понято, скомпилировано и сериализовано компилятором Torch Script. Если вы начинаете с существующей модели PyTorch, написанной с использованием ванильного «нетерпеливого» API, вы должны сначала преобразовать модель в сценарий Torch. В наиболее распространенном случае (описанном ниже) для этого требуется совсем немного усилий. Если у вас уже есть модуль сценариев Torch, вы можете перейти к следующей части этого руководства.

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

Совет: Вы можете найти его в официальномСправочник по скрипту факелаНайдите полную документацию по обоим методам, а также дополнительные рекомендации по их использованию.

Способ 1: преобразовать в сценарий Torch путем трассировки

Чтобы преобразовать модель PyTorch с помощью трассировки в сценарий Torch, экземпляр модели необходимо передать вместе с образцом входных данных вtorch.jit.traceфункция. Это создастtorch.jit.ScriptModuleОбъект, трассировки оценки модели которого будут встроены в модульforwardВ методе:

import torch
import torchvision

# 你模型的一个实例.
model = torchvision.models.resnet18()

# 您通常会提供给模型的forward()方法的示例输入。
example = torch.rand(1, 3, 224, 224)

# 使用`torch.jit.trace `来通过跟踪生成`torch.jit.ScriptModule`
traced_script_module = torch.jit.trace(model, example)

теперь можно отслеживатьScriptModuleОцените его как обычный модуль PyTorch:

In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
In[2]: output[0, :5]
Out[2]: tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

Способ 2: преобразовать в Torch Script через комментарии

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

import torch

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

Поскольку прямой метод этого модуля использует поток управления, зависящий от ввода, он не подходит для отслеживания. Вместо этого мы можем преобразовать его вScriptModule. Чтобы преобразовать модуль вScriptModule, Необходимо использоватьtorch.jit.scriptСкомпилируйте модуль следующим образом:

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)

Если вам нужноnn.Moduleнекоторые методы исключены, поскольку они используютTorchScriptФункции Python, которые еще не поддерживаются, вы можете использовать@torch.jit.ignoreаннотировать это

my_moduleдаScriptModuleЭкземпляр , который можно сериализовать.

Шаг 2: сериализуйте модуль скрипта в файл

Когда у вас есть ScriptModule (путем трассировки или аннотирования модели PyTorch), вы можете сериализовать его в файл. Позже вы сможете загрузить модуль из этого файла с помощью C++ и выполнить его, не полагаясь на Python. Предположим, мы хотим сериализовать показанный ранее пример трассировки.ResNet18Модель. Чтобы выполнить эту сериализацию, просто вызовите модульsaveИ просто передайте имя файла:

traced_script_module.save("traced_resnet_model.pt")

Это сгенерирует в вашем рабочем каталогеtraced_resnet_model.ptдокумент. Если вы также хотите сериализоватьmy_module, пожалуйста позвониmy_module.save("my_module_model.pt")Теперь мы официально покинули область Python и готовы шагнуть в область C++.

Шаг 3: Загрузите модуль скрипта на C++

Чтобы загружать сериализованные модели PyTorch в C++, ваше приложение должно зависеть от API PyTorch C++ (также известного как LibTorch). Дистрибутив LibTorch содержит набор общих библиотек, файлов заголовков и профилей сборки CMake. Хотя CMake не является обязательным требованием для использования LibTorch, это рекомендуемый подход, который будет хорошо поддерживаться в будущем. В этом руководстве мы будем использовать CMake и LibTorch для создания минимального приложения C++, которое просто загружает и выполняет сериализованную модель PyTorch.

минимальное приложение C++

Начнем с обсуждения кода, загружающего модуль. Уже будет сделано следующее:

include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }


  torch::jit::script::Module module;
  try {
    // 使用以下命令从文件中反序列化脚本模块: torch::jit::load().
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    std::cerr << "error loading the model\n";
    return -1;
  }

  std::cout << "ok\n";
}

Заголовок содержит все соответствующие включения из библиотеки LibTorch, необходимые для запуска примера. Наше приложение принимает путь к файлу сериализованного PyTorch ScriptModule в качестве единственного аргумента командной строки, а затем используетtorch::jit::load（）Функция продолжает десериализовать модуль, который принимает этот путь к файлу в качестве входных данных. Взамен мы получаемTorch::jit::script::Moduleобъект. Мы обсудим, как это сделать позже.

Зависит от LibTorch и сборки приложения

Предположим, мы сохранили приведенный выше код в файле с именемexample-app.cppв файле. самый маленькийCMakeLists.txtМожет выглядеть так же просто, как:

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 11)

Последнее, что нужно собрать для примера приложения, — это дистрибутив LibTorch. Вы всегда можете скачать его с сайта PyTorch.страница загрузкиПолучите последнюю стабильную версию на . Если вы скачаете и разархивируете последний архив, вы должны получить папку со следующей структурой каталогов:

libtorch/
  bin/
  include/
  lib/
  share/

  • lib/ папка содержит общие библиотеки, с которыми вы должны связать,
  • Папка include/ содержит заголовочные файлы, которые программа должна включить,
  • папка share/ содержит необходимую конфигурацию CMake для включения вышеописанного простогоfind_package(Torch)Заказ.

Совет: в Windows отладочные и выпускные сборки не совместимы с ABI. Если вы планируете собирать проект в режиме отладки, попробуйте отладочную версию LibTorch.

Завершающим этапом является сборка приложения. Для этой цели предполагается, что макет каталога примеров выглядит следующим образом:

example-app/
  CMakeLists.txt
  example-app.cpp

Теперь мы можем запустить следующую команду изexample-app/Соберите приложение в папке:

mkdir build
cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
make

/path/to/libtorchДолжен быть полный путь к распакованному дистрибутиву LibTorch. Если все пойдет хорошо, это будет выглядеть так:

root@4b5a67132e81:/example-app# mkdir build
root@4b5a67132e81:/example-app# cd build
root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-- The C compiler identification is GNU 5.4.0
-- The CXX compiler identification is GNU 5.4.0
-- Check for working C compiler: /usr/bin/cc
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE
-- Configuring done
-- Generating done
-- Build files have been written to: /example-app/build
root@4b5a67132e81:/example-app/build# make
Scanning dependencies of target example-app
[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
[100%] Linking CXX executable example-app
[100%] Built target example-app

Если мы предоставим модель отслеживания ResNet18, которую мы создали ранее, в двоичный файл примера приложенияtraced_resnet_model.ptпуть, вы должны быть вознаграждены дружеским "хорошо". Обратите внимание, что если вы попытаетесь использоватьmy_module_model.ptЕсли вы запустите этот пример, вы получите сообщение об ошибке, что введенная вами форма несовместима.my_module_model.pt1D нужно вместо 4D.

root@4b5a67132e81:/example-app/build# ./example-app <path_to_model>/traced_resnet_model.pt
ok

Шаг 4. Выполните модуль скрипта на C++.

После успешной загрузки сериализованного ResNet18 в C++ осталось выполнить всего несколько строк кода! Давайте добавим эти строки в приложение C++.main()В функции:

// 创建输入向量
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));

// 执行模型并将输出转化为张量
at::Tensor output = module.forward(inputs).toTensor();
std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

Первые две строки задают ввод для нашей модели. мы создаемtorch::jit::IValueвектор из (значений типа type-errastedScript::Moduleметод принимает и возвращает) и добавляет один вход. Чтобы создать входной тензор, мы используемtorch::ones(), что эквивалентно C++ APItorch.ones. Затем мы запускаемscript::Moduleизforwardметод, передав ему созданный нами входной вектор. Взамен мы получаем новый IValue, вызываяtoTensor()Преобразуйте его в тензор.

Совет. Чтобы узнать больше о таких функциях, как torch::ones и PyTorch C++ API в целом, см. их документацию по адресу https://pytorch.org/cppdocs.
PyTorch C++ API обеспечивает почти такой же функциональный паритет, что и Python API, позволяя вам дополнительно манипулировать и обрабатывать тензоры, как в Python.

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

root@4b5a67132e81:/example-app/build# make
Scanning dependencies of target example-app
[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
[100%] Linking CXX executable example-app
[100%] Built target example-app
root@4b5a67132e81:/example-app/build# ./example-app traced_resnet_model.pt
-0.2698 -0.0381  0.4023 -0.3010 -0.0448
[ Variable[CPUFloatType]{1,5} ]

Для справки, предыдущий вывод Python был:

tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

Похоже, хороший матч!

Подсказка: чтобы переместить модель в память графического процессора, вы можете написать model.to(at::kCUDA);. Убедитесь, что входные данные для модели также находятся в памяти CUDA, вызвав tensor.to(at::kCUDA),
Это вернет новый тензор в память CUDA.

Шаг 5. Получите помощь и изучите API

Мы надеемся, что это руководство даст вам общее представление о пути, который проходят модели PyTorch от Python до C++. Используя концепции, описанные в этом руководстве, вы сможете перейти от ванильной, «нетерпеливой» модели PyTorch к скомпилированной на Python.ScriptModule, в сериализованный файл на диске и, чтобы завершить цикл, в исполняемый скрипт: модуль на C++.

Конечно, есть много концепций, которые мы не рассмотрели. Например, вы можете захотеть использовать собственные расширения операторов, реализованные на C++ или CUDA.ScriptModule, и выполните этот пользовательский оператор в ScriptModule, загруженном в чистую производственную среду C++. Хорошая новость: это возможно и хорошо поддерживается! Теперь вы можете просматриватьэта папкапример в , мы скоро предоставим учебник. На данный момент часто могут быть полезны следующие ссылки:

  • Ссылка на скрипт Torch: https://pytorch.org/docs/master/jit.html
  • Документация API PyTorch С++: https://pytorch.org/cppdocs/
  • Документация по API Python для PyTorch: https://pytorch.org/docs/

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

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

Категории