Как сделать аналог Prisma на Fast Style Transfer, CoreML и TensorFlow

Вступление

Основная часть этого туториала взята из блога Prisma Lab и их работы с PyTorch. Однако мы будем использовать TensorFlow для создания моделей, в частности Fast Style Transfer от Логана Энгстрема. Результатом этого урока будет iOS-приложение, которое будет запускать модели TensorFlow при помощи CoreML. Вот репозиторий в GitHub, в котором содержатся все нужные дополнения.

Благодаря чему это стало возможным:

• Стэнфордское исследование
• Fast Style Transfer
• CoreML от Apple (не поддерживает TensorFlow)
• Google выпускает TensorFlow Lite (не поддерживает CoreML)
• Google выпускает поддержку CoreML (неполная поддержка)
• Мы внесли несколько дополнений и создали готовое решение

Инструменты

Модели: мы будем использовать уже обученные модели FST, но так же хорошо сработают и кастомные модели, вам только нужно будет сделать несколько небольших изменений, которые я упомяну. Скачать модели.

Fast Style Transfer

TensorFlow: При использовании FST лучше всего использовать TensorFlow 1.0.0, а при использовании TensorFlow-CoreML — версию 1.1.0 или новее (для этого урока GPU не понадобится).

TensorFlow CoreML

iOS: 11

Xcode: 9

Python: 2.7

Подготовка

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

Шаг 1: Нам нужно придумать название выходного изображения, иначе TensorFlow сгенерирует его автоматически. Мы можем сделать это в скрипте evaluate.py.

https://gist.github.com/mdramos/729e679446b30b9ca4e3bc56eba884a1#file-eval_node_name-py

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

$ python evaluate.py --checkpoint wave.ckpt --in-path inputs/ --out-path outputs/

> Tensor(“add_37:0”, shape=(20, 720, 884, 3), dtype=float32, device=/device:GPU:0)

Здесь имеет значение только название выходного узла, то есть add_37. Это имеет смысл, так как последний неназванный оператор в сети — это сложение.

Шаг 2: Нам нужно внести ещё несколько изменений в evaluate.py, чтобы сохранить изображение на диск. Обратите внимание, что если вы используете ваши собственные модели, то вам нужно будет добавить код работы с каталогом контрольных точек по сравнению с одним файлом контрольной точки.

https://gist.github.com/mdramos/772ec81683e715f629f1793c75a3d4b1#file-save_graph-py

Шаг 3: Теперь мы запустим evaluate.py для модели и увидим, что ваш файл с изображением сохранен. При обучении моделей вы, вероятно, использовали больше одного образца для тренировки (batch size), а также GPU, однако CoreML принимает только граф с размером 1  и оптимизацией для CPU. Обратите внимание на команду оценки для настройки.

$ python evaluate.py --checkpoint wave/wave.ckpt --in-path inputs/ --out-path outputs/ --device “/cpu:0” --batch-size 1

Отлично, мы создали output_graph.pb и можем приступать к конверсии в CoreML.

Конверсия в CoreML

Благодаря Google теперь существует конвертер из TensorFlow в CoreML. Это отлично, но решение новое, и в нем не хватает некоторых важных операций TensorFlow, например, power.

Шаг 1: Наша модель не будет конвертироваться без поддержки power, но, к счастью, инструменты CoreML предоставляют унарную конверсию, которая поддерживает power. Нам нужно добавить этот код в решение на TensorFlow.

https://gist.github.com/mdramos/4d0900fcd3dbf5ef6c253cc23cf113f3#file-tfcoreml_convert-py

Шаг 2: Создайте и запустите скрипт конверсии.

https://gist.github.com/mdramos/886580af9ebb7939c2457b0ed46dd60d#file-convert-py

$ python convert.py

Актуальный конвертер CoreML не предоставляет возможности вывода изображений из модели. Изображения представлены массивами NumPy (многомерные массивы), которые являются фактическим результатом и скомпилированы для нестандартного типа MultiArray в Swift. Я поискал в интернете и смог получить код, который оценивал выходные данные от CoreML, а затем преобразовывал их в изображения.

Шаг 3: создайте и запустите скрипт трансформации выходных данных в модели my_model.mlmodel.

https://gist.github.com/mdramos/5ae8f76d07356c86cf6b15eda083aba8#file-output-py

$ python output.py

И… 🎉💥🤙… у нас есть рабочая модель CoreML. Я изменил её название на wave в выходном скрипте, а также размер изображения, чтобы тот соответствовал входным данным.

Приложение iOS

Это не совсем урок, поэтому вы можете поработать с моим репозиторием. Я затрону здесь только самое важное.

Шаг 1: Импортируйте модели в свой проект Xcode.

Шаг 2: После импортирования вы сможете инициализировать свои модели:

https://gist.github.com/mdramos/74b96dc98bfe33f4299029ed9026c749#file-init_models-swift

Шаг 3: Создайте класс для входного параметра модели, MLFeatureProvider. img_placeholder — это входные данные, которые определены в скрипте оценки.

https://gist.github.com/mdramos/74b96dc98bfe33f4299029ed9026c749#file-init_models-swift

Шаг 4: Теперь мы можем вызвать модель в нашем коде.

https://gist.github.com/mdramos/75a703bb20176cc2ad2b6acd68ab5ccf#file-code_example-swift

Финал: Остальная часть приложения — это настройки и обработка изображений. Ничего нового или связанного с CoreML, поэтому мы не будем это рассматривать. На этой точке у вас будет понимание того, как все работает вместе, и вы сможете вносить свои улучшения. Я думаю, что как минимум можно улучшить граф FST. Мы должны убрать слишком сложные операции, чтобы финальное приложение работало ещё быстрее. Но и сейчас все работает довольно неплохо.