REST API для инференса на Node.js

Установка и подготовка окружения

Для работы с TensorFlow.js в Node.js необходимо установить основные зависимости:

npm init -y
npm install @tensorflow/tfjs-node express body-parser

@tensorflow/tfjs-node — основной пакет для выполнения моделей TensorFlow в Node.js с ускорением через C++ и CUDA (при наличии GPU).
express — фреймворк для создания REST API.
body-parser — парсер входящих JSON-запросов.

Создание базового сервера выполняется следующим образом:

const express = require('express');
const bodyParser = require('body-parser');
const tf = require('@tensorflow/tfjs-node');

const app = express();
app.use(bodyParser.json());
const PORT = 3000;

Загрузка модели

TensorFlow.js позволяет загружать модели в формате SavedModel или в формате, конвертированном из Python (JSON + веса). Для Node.js используется метод tf.loadLayersModel или tf.node.loadSavedModel:

let model;

async function loadModel() {
    model = await tf.node.loadSavedModel('./model');
}

loadModel().then(() => console.log('Модель загружена'));

Важно: tf.node.loadSavedModel поддерживает нативный формат TensorFlow и обеспечивает более высокую производительность по сравнению с загрузкой через tf.loadLayersModel.

Создание эндпоинта инференса

Эндпоинт REST API принимает входные данные в формате JSON и возвращает предсказания. Пример для модели классификации изображений:

app.post('/predict', async (req, res) => {
    try {
        const { data } = req.body;

        // Преобразование входных данных в тензор
        const inputTensor = tf.tensor(data, [1, 224, 224, 3]); // пример для изображения 224x224 RGB

        // Выполнение предсказания
        const prediction = model.predict(inputTensor);

        // Получение результата в виде массива
        const result = prediction.arraySync();
        res.json({ prediction: result });
    } catch (error) {
        console.error(error);
        res.status(500).json({ error: 'Ошибка при предсказании' });
    }
});

tensor(data, [1, 224, 224, 3]) — создание тензора с указанной формой.
arraySync() — синхронное получение данных из тензора. Для больших данных предпочтительно использовать await prediction.array().

Обработка изображений и данных

Для работы с изображениями в Node.js часто применяются библиотеки sharp или jimp. Пример предобработки изображения:

const sharp = require('sharp');
const fs = require('fs');

async function preprocessImage(path) {
    const buffer = await sharp(path)
        .resize(224, 224)
        .toFormat('png')
        .raw()
        .toBuffer();

    const tensor = tf.tensor(new Uint8Array(buffer), [1, 224, 224, 3]);
    return tensor.div(255.0); // нормализация
}

Нормализация div(255.0) необходима, если модель обучалась на данных с диапазоном [0,1].
Функция возвращает тензор готовый для передачи в model.predict.

Асинхронная обработка запросов

Node.js поддерживает высокую нагрузку через асинхронные вызовы. Для REST API важно избегать блокировки event loop:

app.post('/predict', async (req, res) => {
    try {
        const tensor = await preprocessImage(req.body.imagePath);
        const prediction = await model.predict(tensor).array();
        res.json({ prediction });
    } catch (err) {
        res.status(500).json({ error: err.message });
    }
});

Использование await при загрузке и обработке данных предотвращает блокировку сервера.
Для больших моделей и высокой частоты запросов можно применить пул тензоров через tf.tidy для автоматического освобождения памяти:

const prediction = tf.tidy(() => {
    const input = tf.tensor(data, [1, 224, 224, 3]);
    return model.predict(input);
});

Логирование и мониторинг

Для продакшен-сервера рекомендуется вести журнал запросов и ошибок:

app.use((req, res, next) => {
    console.log(`${req.method} ${req.url}`);
    next();
});

app.use((err, req, res, next) => {
    console.error(err.stack);
    res.status(500).send('Internal Server Error');
});

Логирование помогает выявлять проблемные запросы и аномалии в работе модели.
Для крупных систем можно интегрировать winston или pino для структурированных логов.

Запуск сервера

Старт сервера выполняется стандартно:

app.listen(PORT, () => {
    console.log(`Server running on port ${PORT}`);
});

При изменении кода удобно использовать nodemon для автоматической перезагрузки.
При работе с GPU и тяжелыми моделями важно контролировать потребление памяти и производительность через tf.engine().memory().