Прогнозування на стороні клієнта за допомогою TensorFlow.js

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

Сьогодні я поділюся з вами своїм досвідом і покажу, як розгорнути модель у такий спосіб, щоб частина розрахунків відбувалася на стороні клієнта. Ця стаття призначена для всіх, хто створив модель і бажає зменшити навантаження на сервер, передавши частину з прогнозуванням клієнтові. Особливо для Data Scientist-ів, які використовують Python щодня і погано володіють Javascript.

Вступ

Уявімо, що ви створили якусь чудову модель, яка робить круті речі і допомагає людям. Наприклад, модель прогнозує улюблений емоджі людини на основі фотографії її чашки. Ви завантажили цю модель в інтернет. Щоденне використання сягає приблизно 1000 запитів — небагато. Простий сервер може дати з цим раду, але одного дня про цю модель дізнається багато людей, і ви почнете отримувати по 100 тисяч запитів щоденно. Ваш сервер, швидше за все, «помре». Отже, ви можете або збільшити сервер і додавати щоразу більше пам’яті, або переписати прогнозування на сторону клієнта. Якщо ви виберете другий варіант, то ось для вас туторіал.

Щоб досягти мети, нам потрібні такі компоненти:

• Backend: Flask, будь-яка бібліотека для попередньої обробки зображення в Python.
• Frontend: TensorFlow.js

Нещодавно в TensorFlow.js з’явилася підтримка Node.js, проте ми будемо використовувати Flask, який є бібліотекою Python. Часто деяким натренованим моделям потрібна попередня обробка даних для коректної роботи. Наразі попередню обробку набагато зручніше виконувати в Python ніж в JavaScipt. Сподіваюся, що одного разу стане також можливою і попередня обробка на стороні клієнта.

Створення моделі

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

Після того, як модель збереглася, у вас з’явиться папка з таким вмістом.

Де group\*-shard\*of\* — це колекція бінарних weight файлів і model.json — це модель топології і конфігурація.

Налаштування Flask-сервера

Нам потрібно, щоб користувачі мали доступ до нашої моделі.

Поки що сервер дуже простий, нічого складного, лише один маршрут, що повертає сторінку index.html. Загальна структура системи виглядає ось так.

Створення index.html

Нам потрібна точка входу, з якою користувач буде взаємодіяти і де буде відбуватися наше прогнозування. Отож, давайте налаштуємо index.html.

Наша перша версія виглядатиме так. Єдина важлива річ тут — на рядку 6, де ми додаємо Tensorflow.js з CDN. Наступний крок — це додавання тіла HTML, щоб користувач зміг завантажувати зображення і натискати на кнопки :) Ось він.

Наступний і останній крок для частини з HTML — додати трохи стилю до нашої сторінки, відповідно присвоївши класи елементам HTML, і також створити основний main.js файл який буде містити наше магічне прогнозування. Тепер давайте глянемо на остаточну версію index.html.

Ваш index.html може відрізнятися від мого, можете додати або ж видалити деякі його частини. Тим не менше, найважливіші тут:

• рядок 6 (додає скрипт з CDN в head — погана практика для продакшн-версії коду, проте я не буду пояснювати, що таке npm і node_modules);
• рядок 13 (ввід, щоб користувач мав змогу завантажити зображення, можете використовувати різні типи вводу);
• рядок 20 (наш скрипт, що відповідає за прогнозування на стороні клієнта).

Створення main.js

Настав час застосовувати магію. Передусім нам потрібно ініціалізувати кнопки, ввід, модель і функцію для прогнозування.

В рядку 3 — адреса нашої моделі, зараз вона знаходиться на моїй локальній машині, але взагалі її можна розгорнути будь-де. Також пізніше ми створимо маршрут у Flask для цієї моделі.

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

Ми надсилаємо зображення на /api/prepare/, цей шлях ми додамо пізніше. Також ми присвоюємо відповіді сервера на поле із зображенням tf.tensor2d.

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

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

Ви можете переписати частину з fetch так, щоб надіслати усі файли в одному пості. Не забудьте привести повернений масив до того ж формату, за яким тренувалася модель.

Оновлення сервера Flask

Тепер нам потрібно оновити наш сервер, щоб він міг виконувати попередню обробку зображення для /api/prepare/ і також виводити модель для /model на frontend. Фінальна версія сервера виглядатиме приблизно так.

Для попередньої обробки у нас є дві функції:

• prepare (викликається на /api/prepare/);
• preprocessing (бере зображення й повертає змінене зображення до масиву numpy, ця функція може виконувати будь-яку попередню обробку, і все працюватиме без проблем, якщо вона повертає масив numpy).

Модель:

• model (викликається для /model);
• load_shards (викликається для будь-якого файлу, який можна викликати з /, ця функція використовується для завантаження бінарних weight файлів).

Чому нам потрібні дві функції та два окремі API для моделі замість одного?

В цій версії TensorFlow.js, коли ми завантажуємо модель для якогось API,

model = await tf.loadModel(modelURL);

вона спочатку завантажує модель, яка є файлом JSON, з modelURL, і після цього автоматично надсилає ще декілька POST запитів до основи домену, щоб завантажити shards (перегляньте цей запит в демо, у логах сервера). Оскільки я не хочу тримати модель в основі шляху разом із сервером, мені потрібні дві функції: одна для model.json та інша для shards.

Результат

На цьому все! Тепер ви можете передати клієнту процес прогнозування улюбленого емоджі людини на основі зображення її чашки або, як це у моєму випадку, MNIST! Якщо все зроблено правильно, то ви побачите щось таке.

Дякую що прочитали! Можете додати/переписати будь-яку частину, як вам завгодно. Насолоджуйтесь!

Цю статтю ви також можете прочитати англійською мовою.