[Перевод] Альтернатива ML-Agents: интегрируем нейросети в Unity-проект с помощью PyTorch C++ API

3su_p-oszzalsgau9mfba9xllhs.png

Кратко объясню, что будет происходить в этой статье:

  • покажу, как использовать PyTorch C++ API для интеграции нейросети в проект на движке Unity;
  • сам проект я подробно описывать не буду, это не имеет значения для данной статьи;
  • использую готовую модель нейросети, преобразовав её трассировку в бинарник, который будет подгружаться в рантайме;
  • покажу, что такой подход существенно облегчает деплой сложных проектов (например, нет проблем с синхронизацией сред Unity и Python).

Добро пожаловать в реальный мир


Методы машинного обучения, в том числе нейронные сети, по-прежнему очень комфортно себя чувствуют в экспериментальных окружениях, а запуск таких проектов в реальном мире часто вызывает трудности. Я немного расскажу про эти трудности, опишу ограничения по способам выхода из них, а также дам поэтапное решение проблемы интеграции нейросети в Unity-проект. 

Другими словами, мне нужно превратить исследовательский проект на PyTorch в готовое решение, способное вместе с движком Unity работать в боевых условиях.
Можно несколькими способами интегрировать нейронную сеть в Unity. Я предлагаю использовать C++ API для PyTorch (под названием libtorch) для создания нативной разделяемой библиотеки, которую затем можно будет подключить к Unity как плагин. Существуют и другие подходы (например, использовать ML-Agents), которые в определённых случаях могут быть проще и эффективнее. Но преимущество моего подхода заключается в том, что он обеспечивает большую гибкость и даёт больше возможностей. 

Допустим, у вас есть какая-то экзотическая модель и вы просто хотите использовать существующий PyTorch-код (который был написан без намерения общаться с Unity); или ваша команда разрабатывает собственную модель и не хочет отвлекаться на мысли о Unity. В обоих случаях код модели может быть сколь угодно сложным и использовать все возможности PyTorch. А если вдруг дело дойдёт до интеграции, в игру вступит C++ API и завернёт всё в библиотеку без малейшего изменения изначального PyTorch-кода модели.

Итак, мой подход сводится к четырём ключевым шагам:

  1. Настройка окружения.
  2. Подготовка нативной библиотеки (C++).
  3. Импорт функций из библиотеки / подключение плагина (Unity / C#). 
  4. Сохранение / развёртывание модели.

ВАЖНО: поскольку я делал проект, сидя под Linux, некоторые команды и настройки отталкиваются от этой ОС;, но не думаю, что здесь что-либо должно слишком зависеть от неё. Поэтому вряд ли подготовка библиотеки под Windows вызовет трудности.

Настройка окружения


Прежде чем устанавливать libtorch, убедитесь, что у вас есть
А если хотите использовать GPU, вам потребуются:
  • CUDA toolkit (на момент написания статьи была актуальна версия 10.1);
  • CUDNN library

С CUDA могут возникнуть сложности, потому что драйвер, библиотеки и прочая хурма должны дружить между собой. И вам придётся поставлять эти библиотеки вместе с Unity-проектом — чтобы всё работало из коробки. Так что для меня это самая неудобная часть. Если вы не планируете использовать GPU и CUDA, то знайте: вычисления замедлятся в 50–100 раз. И даже если у пользователя довольно слабый графический процессор — лучше с ним, чем без него. Даже если ваша нейросеть включается в работу довольно редко, эти редкие включения приведут к задержке, которая будет раздражать пользователя. Возможно, в вашем случае всё будет иначе, но… нужен ли вам этот риск?

После того, как вы установили означенное выше ПО, пора загрузить и (локально) установить libtorch. Необязательно устанавливать для всех пользователей: можно просто поместить её в каталог своего проекта и обратиться к нему при запуске CMake.

Подготовка нативной библиотеки


Следующий шаг — конфигурирование CMake. Я взял за основу пример из документации PyTorch и изменил его так, чтобы после сборки мы получали библиотеку, а не исполняемый файл. Положите этот файл в корневой каталог вашего проекта с нативной библиотекой.CMakeLists.txt
cmake_minimum_required(VERSION 3.0 FATAL_ERROR)

project(networks)

find_package(Torch REQUIRED)

set(CMAKE_CXX_FLAGS «${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}»)

add_library(networks SHARED networks.cpp)

target_link_libraries(networks «${TORCH_LIBRARIES}»)

set_property(TARGET networks PROPERTY CXX_STANDARD 14)

if (MSVC)

	file(GLOB TORCH_DLLS «${TORCH_INSTALL_PREFIX}/lib/*.dll»)

	add_custom_command(TARGET networks

		POST_BUILD

		COMMAND ${CMAKE_COMMAND} -E copy_if_different

		${TORCH_DLLS}

		$)

endif (MSVC)

Исходный код библиотеки будет размещён в networks.cpp

В этом подходе есть ещё одна приятная особенность: нам пока не нужно думать, какую именно нейронную сеть мы хотим использовать с Unity. Причина (немного забегая вперед) заключается в том, что мы в любой момент можем запустить сеть в Python, получить её трассировку и просто сказать libtorch «применить эту трассировку для этих входов». Поэтому, можно сказать, что наша нативная библиотека просто обслуживает некий чёрный ящик, работая с вводом-выводом.

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

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

networks.cpp
#include 

#include 

#include  

extern «C»

{

// This is going to store the loaded network

torch::jit::script::Module network;

Чтобы вызывать функции нашей библиотеки непосредственно из Unity, нужно передать информацию об их точках входа. В Linux я использую для этого __attribute__((visibility («default»))). В Windows для этого существует спецификатор __declspec (dllexport), но, честно говоря, я не проверял, работает ли он там.

Итак, начнём с функции загрузки трассировки нейросети с диска. Файл имеет относительный путь — он лежит в корневом каталоге проекта Unity, а не в Assets/. Так что будьте внимательны. Вы также можете просто передать имя файла из Unity.


extern __attribute__((visibility(«default»))) void InitNetwork()

{
	network = torch::jit::load(«network_trace.pt»);

	network.to(at::kCUDA); // If we're doing this on GPU
}

Теперь перейдём к функции, которая кормит сеть входными данными. Напишем на С++ код, который использует указатели (ими управляет Unity) для перегонки данных туда и обратно. В этом примере я полагаю, что моя сеть имеет входы и выходы фиксированной размерности и запрещаю Unity менять это. Здесь, например, я возьму Tensor {1,3,64,64} и Tensor {1,5,64,64} (например, такая сеть нужна для сегментации пикселей RGB-изображений на 5 групп).

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

Чтобы преобразовать данные в формат, с которым работает libtorch, мы используем функцию torch: from_blob. Она принимает массив чисел с плавающей запятой и описание тензора (с указанием размерности) и возвращает созданный Тензор.

Нейросети могут принимать несколько входных аргументов (например, вызов forward () принимает x, y, z в качестве входных данных). Чтобы справиться с этим, все входные тензоры упаковываются в вектор стандартной библиотеки шаблонов torch: jit: IValue (даже если аргумент только один).

Чтобы получить данные из тензора, проще всего обработать их поэлементно, но если из-за этого упадёт скорость обработки, для оптимизации процесса чтения данных можно использовать Tensor: accessor. Хотя лично мне это не понадобилось.

В результате для моей нейросети получается вот такой простой код:

extern __attribute__((visibility(«default»))) void ApplyNetwork(float *data, float *output)

{

Tensor x = torch::from_blob(data, {1,3,64,64}).cuda();

std::vector inputs;

inputs.push_back(x);

Tensor z = network.forward(inputs).toTensor();

for (int i=0;i<1*5*64*64;i++)

output[i] = z[0][i].item();

}

}

Чтобы скомпилировать код, следуйте указаниям в документации, создайте подкаталог build/ и запустите следующие команды:
cmake -DCMAKE_PREFIX_PATH=/absolute/path/to/libtorch ..

cmake --build . --config Release

Если всё пойдёт хорошо, будут созданы файлы libnetworks.so или networks.dll, которые вы сможете разместить в Assets/Plugins/ вашего Unity-проекта.

Подключение плагина к Unity


Для импорта функций из библиотеки используем DllImport. Первая функция, которая нам понадобится, — это InitNetwork (). При подключении плагина Unity вызовет именно её:
using System.Runtime.InteropServices;

public class Startup : MonoBehaviour

{

...

[DllImport(«networks»)]

private static extern void InitNetwork();

void Start()

{

...

InitNetwork();

...

}

}

Чтобы движок Unity (С#) мог обмениваться данными с библиотекой (C++), я поручу ему всю работу по управлению памятью:
  • выделю память под массивы нужного размера на стороне Unity;
  • передам адрес первого элемента массива в функцию ApplyNetwork (её тоже перед этим нужно импортировать);
  • просто позволю адресной арифметике C++ обращаться к этой памяти при получении или отправке данных.

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

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

[DllImport(«networks»)]

private static extern void ApplyNetwork(ref float data, ref float output);

void SomeFunction() {

float[] input = new float[1*3*64*64];

float[] output = new float[1*5*64*64];

// Load input with whatever data you want

...

ApplyNetwork(ref input[0], ref output[0]);

// Do whatever you want with the output

...

}

Сохранение модели


Статья близится к концу, а мы так и обсудили, какую нейронную сеть я выбрал для моего проекта. Это простая свёрточная нейросеть, которую можно использовать для сегментации изображений. Я не включил в модель сбор данных и обучение: моя задача — рассказать про интеграцию с Unity, а не про заморочки с трассировкой сложных нейросетей. Не обессудьте.

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

Так может выглядеть Python-код для нашей простой модели:

import torch

import torch.nn as nn

import torch.nn.functional as F

class Net(nn.Module):

def __init__(self):

super().__init__()

self.c1 = nn.Conv2d(3,64,5,padding=2)

self.c2 = nn.Conv2d(64,5,5,padding=2)

def forward(self, x): z = F.leaky_relu(self.c1(x)) z = F.log_softmax(self.c2(z), dim=1)

return z 

Код не очень красивый, конечно, но, думаю, идея понятна.

Сохранить (экспортировать) модель с текущими значениями коэффициентов можно так:

network = Net().cuda()

example = torch.rand(1, 3, 32, 32).cuda()

traced_network = torch.jit.trace(network, example)

traced_network.save(«network_trace.pt»)

Развёртывание модели


Мы сделали статическую библиотеку, но для развёртывания этого недостаточно: в проект нужно включить дополнительные библиотеки. К сожалению, у меня нет 100-процентной уверенности в том, какие именно библиотеки нужно включить обязательно. Я выбрал libtorch, libc10, libc10_cuda, libnvToolsExt и libcudart. В сумме они добавляют 2 Гб к изначальному размеру проекта. 

LibTorch vs ML-Agents


Я считаю, что для многих проектов, особенно в области исследований и прототипирования, действительно стоит выбрать ML-Agents, плагин, созданный специально для Unity. Но когда проекты становятся более сложными, нужно подстраховаться — на случай, если что-то пойдёт не так. А такое случается нередко…

Пару недель назад я как раз использовал ML-Agents для взаимодействия между демо-игрой на Unity и парой нейронных сетей, написанных на Python. В зависимости от игровой логики Unity вызывал одну из этих сетей с разными наборами данных.

Мне пришлось основательно покопаться в Python API для ML-Agents. Некоторые операции, которые я использовал в моих нейросетях, например 1d свёртка и операции транспонирования, не поддерживались в Barracuda (это библиотека трассировки, которую в настоящее время использует ML-Agents).

Проблема, с которой я столкнулся, заключалась в том, что ML-Agents собирает «запросы» от агентов в течение некого временного интервала, а затем для оценки отправляет, к примеру, в Jupyter notebook. Однако некоторые из моих нейросетей зависели от результатов работы других моих сетей. И, чтобы получить оценку всей цепочки моих нейронных сетей, мне пришлось бы каждый раз, делая запрос, ждать какое-то время, получать результат, делать другой запрос, ждать, получать результат и так далее. Кроме того, порядок включения этих сетей в работу нетривиально зависел от ввода данных пользователем. А это означало, что я не мог просто последовательно запускать нейросети. 

Кроме того, в некоторых случаях объём данных, который мне нужно было отправить, должен был варьироваться. А ML-Agents больше рассчитан на фиксированную размерность для каждого агента (вроде бы её можно менять на лету, но я отношусь к этому скептически).

Я мог бы сделать что-то вроде вычисления последовательности вызова нейросетей по запросу, отправляя соответствующие входные данные для Python API. Но из-за этого мой код, как на стороне Unity, так и на стороне Python, стал бы слишком сложным, или вовсе избыточным. Поэтому я решил изучить подход с использованием libtorch и не прогадал.

Если бы раньше кто-нибудь попросил меня встроить в Unity-проект предсказательную модель GPT-2 или MAML, я бы посоветовал ему постараться обойтись без этого. Реализация такой задачи с использованием ML-Agents слишком сложна. Но теперь я могу найти или разработать любую модель с PyTorch, а потом завернуть её в нативную библиотеку, которая подключается к Unity как обычный плагин.


Облачные серверы от Маклауд быстрые и безопасные.

Зарегистрируйтесь по ссылке выше или кликнув на баннер и получите 10% скидку на первый месяц аренды сервера любой конфигурации!

et1aypandyuamqprsz3m2ntm4ky.png

© Habrahabr.ru