Від ML-моделей до генеративних мереж: як ми створюємо ШІ-ефекти

💡 Усі статті, обговорення, новини про AI — в одному місці. Приєднуйтесь до AI спільноти!

Усім привіт! Мене звати Андрій Рижий. Я — МL Lead Engineer в українській продуктовій компанії Skylum, і останні п’ять років працюю над створенням ШІ-ефектів. У цьому матеріалі хочу розказати про специфіку створення ШІ-ефектів для фоторедактора Luminar Neo.

У своїй щоденній роботі ми використовуємо комбінації алгоритмів класичного комп’ютерного зору і підходів Deep Learning, які навчаємо під певні конкретні задачі. Далі у статті я поділюся підходами, кейсами та принципами роботи нашої команди, що дозволяють створювати якісні та популярні ШІ-ефекти.

Наші найкращі ефекти — це комбінація класичних алгоритмів і ML-підходів

Попри очевидні переваги ML, поки що мало моделей працюють фотореалістично та з фотографіями високої роздільної здатності. Саме комбінація з класичними алгоритмами допомагає ML-моделям краще розуміти контекст, а отже — генерувати більш реалістичні й узгоджені з контекстом результати.

Такий комбінований підхід ми використали, зокрема, в Studio Light Tool (створення студійного освітлення для фото) та Sky AI (заміна неба). Наприклад, в ефекті Studio Light Tool завдяки ML-моделі ми знаходимо людину, лице і глибину обличчя людини на фотографії. Тобто, будуємо 3D-об’єкт.

Потім, спираючись на всі ці дані, додаємо класичні алгоритми: це дає можливість розмістити світло у будь-якій точці та переміщувати його під потрібним кутом. А оскільки маємо математичну модель обличчя, отриману за допомогою нейронних мереж, то здатні вираховувати зміну кольору у будь-якій точці обличчя.

Додаю приклад, як саме виглядає глибина обличчя та, власне, Studio Light Tool.

А в роботі зі Sky AI ми використовуємо нейромережу для сегментації різних об’єктів на фотографії, такі як небо, воду, відображення неба у воді. Також окремо визначаємо лінію горизонту. Якщо на фото присутня людина, то сегментуємо її також. Ці маски дають змогу фотореалістично замінити небо та правильно відтворити колір на кожній точці не-неба.

Приміром, ось так виглядає приклад сегментації неба, води, відображення неба у воді, і власне, ефекту Sky AI.

Збір даних — один з ключових етапів у створенні моделей ШІ-ефектів

Зазвичай ми використовуємо дані умовно двох типів. По-перше, це дані для задач, які передбачають пошук об’єктів. Приміром, зображення і маски, зображення і bounding boxes, зображення і клас, до якого належить дане зображення.

По-друге, це дані у форматі «до-після», які працюють на ML-основі. Залежно від задачі ми чітко розуміємо, які дані подавати на вхід моделі (зазвичай це фотографії або їх частини) та який результат очікувати на виході (покращене чи збільшене зображення). З урахуванням цього формується датасет, що складається з пар «вхід-вихід» та, відповідно, використовується для тренування моделі.

Утім, частіше за все дані, потрібні нашим моделям для навчання, у відкритому доступі знайти доволі важко. Тому з’являється необхідність створювати та розмічати фотографії власними силами. У нашій практиці був кейс з dust spots — дефектом, що виникає при попаданні пилу на матрицю фотоапарата під час зміни об’єктива. Ця проблема проявляється у вигляді дратівливої точки на фото.

Виявилося, що знайти у мережі достатню кількість фото з цим дефектом для тренування моделі просто неможливо. А для навчання нам потрібні були тисячі таких фотографій. Тож вирішили звернутися до наших користувачів і попросили їх надіслати чимбільше світлин з даст-спотами. У підсумку, у нас накопичилось достатньо матеріалу і ми успішно використали його в роботі.

Для наочності додаю приклад, як виглядають dust spots. Вони можуть бути і у вигляді простого кола, і складнішої форми.

І ось приклад, як ми їх знаходимо:

Оскільки самі споти доволі малі, тому просто використовувати сегментаційну модель не можна. При поданні великого розширення у сегментаційну модель дуже сильно страждає перфоманс. А при малому розширенні (наприклад, 512×512) — спотів майже непомітно, і відповідно модель їх не знаходить.

З цієї ситуації ми знайшли такий вихід: навчання окремого детектора малих об’єктів для знаходження локалізації спота і окрема невелика модель сегментації, яка сегментує кожен окремий спот.

Ось приклад псевдокоду, який це виконує сам інференс:

detector_model = load_detector_model()
segmentation_model = load_segmentation_model()
detector_input = detector_preprocessing(input_image)
detector_output = detector_postprocessing(detector_model.predict(detector_input))
dust_segmentation = np.zeros_like((input_image.shape[0], input_image.shape[1]))
for detection in detector_output:
input_crop = input_image[detection[0]:detection[1], detection[2]:detection[3], :]
segmentation_crop = segmentation_model.predict(segmentation_preprocessing(input_crop))
segmentation_crop = segmentation_crop > 0.5
dust_segmentation[detection[0]:detection[1], detection[2]:detection[3]] = segmentation_crop

Як детектор використовуємо одну з варіацій yolo (ultralytics), яку навчено з деякими трюками для знаходження невеликих об’єктів. У якості швидкого сегментатора використовуємо MobilenetV2 у комбінації з deeplabv3.

За відсутності потрібних зображень у загальнодоступних джерелах, також їх доцільно генерувати штучно. Так, наприклад, для розробки моделі для знаходження глибини ми використовували ігровий рушій, де за допомогою додаткових налаштувань можна робити спеціальні скриншоти, в яких буде зберігатись пара зображень — вхідне зображення та карта глибини, яку можна використовувати як істину для навчання моделі.

Для знаходження глибини ми використовуємо Unreal Engine, в якому штучно генеруємо датасет. Ось приклади:

Головне про цикл розробки кожної моделі ШІ-ефекту

Спочатку потрібно розмітити деяку кількість даних і повчити модель. На цьому етапі варто подивитися, наскільки модель гарно працює на кейсах, близьких до тих, з якими зіштовхуватимуться користувачі. Далі варто знайти бед кейси, розмітити їх та, відповідно, доповнити датасет. Після цього залишається перевчити моделі та усе протестувати.

Ми використовуємо цю послідовність дій доти, поки не отримуємо гарних результатів. У випадках, коли є багато публічних даних, які можна використовувати згідно з ліцензією, розробка одного ефекту триває зазвичай один-два місяці. У такому разі основна частина роботи будується навколо підбору архітектури та навчання самої моделі.

Бувають і випадки, коли розробка ефекту займає близько року. З ефектом боке, наприклад, у нас цикл розробки включив 12 ітерацій і тривав набагато більше часу, ніж очікувала вся команда. Серед іншого, складність полягала в тому, щоб поцілити у правильні дані, які відображають вибірку наших користувачів.

Спочатку ми тренували модель на фотографіях людей. Однак вона не могла правильно обробляти світлини людей з рюкзаками, букетами, скейтами, велосипедами тощо. Знадобилося дюжина сценаріїв з доповненням датасету, зокрема і бед кейсами, щоб модель навчилася обробляти будь-які сценарії на фотознімку коректно.

Приклад роботи ефекту Bokeh AI:

Наведу приклад того, як датасет допоміг покращити сегментацію, яку ми використовували у Bokeh AI. На цих фото результат роботи двох однакових мереж, що були навчені на різних датасетах. Перша — на датасеті, де присутні тільки люди. Друга — на датасеті, де присутні люди і також люди, які користуються транспортом, коли транспорт теж є головним об’єктом (15 % датасету):

Можна побачити, що навіть невелика частина з таргетними даними суттєво впливає на результат.

Варто враховувати складність створення ML-моделей для девайсів користувачів

Створювати МL-моделі, які запускаються на девайсах користувачів, в рази складніше, ніж розробляти їх для серверів. Оскільки ноутбуки мають значно меншу обчислювальну потужність, то виникає потреба у ретельній оптимізації. А саме, в частині зменшення розміру моделі, підвищення швидкості роботи та адаптації моделі до конкретного девайса. При цьому, не повинна страждати якість, чого досягти за таких обставин буває доволі складно.

Для того, аби розв’язати цю проблему, ми тісно співпрацюємо з виробниками комп’ютерних компонентів: Microsoft, Intel, Qualcomm, Nvidia тощо. Ці компанії діляться своїми фреймворками та найновішими розробками, імплементація яких значно покращує продуктивність ML-моделей на персональних комп’ютерах. Йдеться насамперед про фреймворки OpenVINO, SNPE, DirectML та CoreML. Загалом чим новіший комп’ютер, тим Luminar Neo краще і головне — швидше працюватиме на девайсах користувачів.

Ось, наприклад, скрипт конвертації сегментаційної моделі в coreml:

import torch
import coremltools as ct
import numpy as np
from PIL import Image
from torchvision import transforms
import urllib
class Deeplabv3MobilenetV3(torch.nn.Module):
    def __init__(self):
        super(Deeplabv3MobilenetV3, self).__init__()
        self.model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_mobilenet_v3_large', pretrained=True).eval()
    def forward(self, x):
        res = self.model(x)
        x = res["out"]
        return x
model = Deeplabv3MobilenetV3()
url, filename = ("<https://github.com/pytorch/hub/raw/master/images/deeplab1.png>", "deeplab1.png")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)
input_image = Image.open(filename)
input_image = input_image.convert("RGB")
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
example_input = input_tensor.unsqueeze(0)
out = model(example_input)
traced_model = torch.jit.trace(model, example_input)
out = traced_model(example_input)
coreml_model = ct.convert(
    traced_model,
    convert_to="mlprogram",
    inputs=[ct.TensorType(shape=example_input.shape)]
)
coreml_model.save("deeplabv3.mlpackage")
out_coreml = coreml_model.predict({'x':example_input})['var_1055']
print('difference: ', np.square(out_coreml - out.detach().numpy()).mean())

Результат — difference: 9.996132e-05.

Як бачимо, різниця між виходами моделі не дуже велика. Давайте заміряємо швидкодію цих моделей (на девайсі mac book pro 2,6 GHz 6-Core Intel Core i7). Як правило, ми замірюємо час ініту моделі, перший запуск та середнє з наступних запусків. Такі числа дають нам повноцінне порівняння двох моделей:

import os
import torch
import numpy as np
import time
import coremltools as ct
example_input = torch.rand(1, 3, 224, 224)
st = time.time()
model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_mobilenet_v3_large', pretrained=True).eval()
torch_init_time = time.time() - st
torch_times = []
for i in range(11):
    st = time.time()
    model(example_input)
    torch_times.append(time.time() - st)
st = time.time()
model = ct.models.MLModel('deeplabv3.mlpackage')
coreml_init_time = time.time() - st
coreml_times = []
for i in range(11):
    st = time.time()
    model.predict({'x':example_input})
    coreml_times.append(time.time() - st)
print('torch_init_time: ', torch_init_time)
print('torch first run: ', torch_times[1])
print('torch avg 10 times: ', np.mean(torch_times[1:]))
print('coreml_init_time: ', torch_init_time)
print('coreml first run: ', coreml_times[1])
print('coreml avg 10 times: ', np.mean(coreml_times[1:]))

Результат

torch_init_time: 0.2949812412261963

torch first run: 0.1378791332244873

torch avg 10 times: 0.1355872392654419

coreml_init_time: 0.2949812412261963

coreml first run: 0.013335943222045898

coreml avg 10 times: 0.011963438987731934

Можна побачити, що модель в Coreml-фреймворку працює в 10 разів швидше (на нових M-процесорах буде ще швидше), що дає неабиякий приріст у перфомансі продукту. Також варто звернути увагу на розмір самої моделі:

pytorch-версія — 42.5 Mb

coreml-версія — 22.2 Mb

Оскільки coreml за замовчуванням використовує тип даних float16 для зберігання, то і місця на диску модель займає менше.

Такого ж роду оптимізацій можна досягти й за допомогою інших платформних фреймворків:

• OpenVino — має особливі алгоритми квантизації, що дозволяє використовувати модель в типі даних uint8, а отже, зменшувати модель аж до 4 разів. При запуску на останніх CPU, GPU від Intel дає солідний приріст у швидкості;
• SNPE — фреймворк від Qualcomm, що дозволяє моделям працювати на останніх процесорах від Qualcomm arm архітектури. Після оптимізацій і квантування дає пришвидшення аж до 100 разів в порівнянні до звичайного запуску без оптимізацій;
• DirectML — фреймворк від Microsoft дозволяє запускати моделі на всіх GPU, які підтримує Windows. І це суттєво покращує перфоманс в деяких кейсах;
• Tensorrt — фреймворк від Nvidia, він дозволяє оптимізувати і пришвидшувати моделі, які запускаються на продуктах від Nvidia. Його дуже зручно і ефективно використовувати у клауді.

Генеративні мережі — нова ера можливостей

Поява конволюційних мереж у 2017 році спричинила нову хвилю розвитку у сфері ШІ. Завдяки генеративним мережам ми маємо змогу працювати з новим інструментом — дифузійними моделями, які мають значний по тенціал для розв’язання нових задач. Так, тренувати їх виходить набагато довше і дорожче, але з іншої сторони вони дають обнадійливі результати. Останні наші ефекти — GenErase, GenExpand та GenSwap — були розроблені на основі дифузійних моделей. Використання цієї технології дозволило фотореалістично видаляти об’єкти з фото, замінювати їх на інші та розширювати саме зображення.

Генеративні моделі багаті на нові можливості, але не варто забувати, що водночас вони вимагають чимало ресурсів і для процесу навчання, і для інференсу. Через це запускати такі великі моделі на кастомерських девайсах недоцільно. Натомість варто зосередити зусилля навколо оптимізації генеративних моделей на серверах, що дозволить значно зменшити їхню ресурсомісткість.

І в якості підсумку додам кілька важливих порад, на що, на мою думку, потрібно звернути увагу при роботі з ML та DL-моделями

1. Краще створювати власні дані, ніж сподіватися на відкриті джерела. Певна річ, існують публічні датасети, але їхні можливості доволі обмежені. Для того, щоб робити унікальні речі, потрібно працювати з унікальними даними. Можливо є сенс створити власний профільний відділ або найняти фахівців на аутсорсі, які б предметно займалися постачанням якісних даних для проєктів.
2. За пошуку додаткової інформації звертайте увагу на наукові статті лише з викладеною імплементацією, звісно якщо це не статті лідерів індустрії, таких як OpenAI та інших. З мого досвіду, немає сенсу витрачати час на тексти без робочого коду або прототипу, яких за останній час значно побільшало у загальному доступі. Також обов’язково переконайтеся, що викладений у матеріали код відповідає результатам статті.
3. Надавайте перевагу перевірці гіпотез на практиці замість абстрактному теоретизуванню. Якщо маєте хоча б один сумнів щодо запропонованої вами гіпотези, краще одразу її перевірити. Замінити шари у нейронній мережі, її перенавчити та подивитися на результат куди продуктивніше, ніж розмірковувати про різні гіпотетичні сценарії.