Экспорт моделей в TensorFlow.js

ConvNetJS — это библиотека для нейронных сетей на чистом JavaScript, разработанная с целью быстрого прототипирования и обучения небольших сетей прямо в браузере. Несмотря на ограниченный функционал по сравнению с современными фреймворками вроде TensorFlow.js, ConvNetJS позволяет создавать, обучать и сохранять модели, что делает возможным их последующую конвертацию для использования в более мощных системах.

Структура модели ConvNetJS

Модель в ConvNetJS строится на основе объекта net, который содержит последовательность слоев (layers) и веса (params). Каждый слой описывается своим типом (conv, fc, softmax, relu и т.д.) и набором параметров, таких как размер ядра свертки, количество фильтров, шаг свертки и функция активации.

Пример определения сети:

var net = new convnetjs.Net();
var layer_defs = [];
layer_defs.push({type:'input', out_sx:28, out_sy:28, out_depth:1});
layer_defs.push({type:'conv', sx:5, filters:8, stride:1, pad:2, activation:'relu'});
layer_defs.push({type:'pool', sx:2, stride:2});
layer_defs.push({type:'fc', num_neurons:10, activation:'softmax'});
net.makeLayers(layer_defs);

После обучения модель хранит веса каждого слоя в объекте net.layers, доступ к которым осуществляется через методы getParamsAndGrads().

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

ConvNetJS предоставляет встроенные методы сериализации модели в JSON. Сохраняются как структура сети, так и веса.

Пример сохранения:

var json = net.toJSON();
var jsonStr = JSON.stringify(json);
// jsonStr можно сохранить в файл или отправить на сервер

Объект JSON содержит:

  • layers: массив слоев с их типами и параметрами
  • params: массив весов и смещений для каждого слоя
  • version: версия библиотеки

Подготовка к экспорту в TensorFlow.js

TensorFlow.js использует формат TensorFlow SavedModel или tfjs Layers Model (JSON + бинарные веса). Чтобы перенести модель из ConvNetJS, необходимо:

  1. Преобразовать структуру слоев в совместимый формат TensorFlow.js.
  2. Сопоставить функции активации и слои.
  3. Экспортировать веса в массивы TensorFlow.js.

Сопоставление слоев и активаций

ConvNetJS TensorFlow.js (tf.layers)
input tf.input
conv tf.layers.conv2d
pool tf.layers.maxPooling2d
fc tf.layers.dense
relu activation='relu'
softmax activation='softmax'

Важно учитывать порядок весов: ConvNetJS использует [filter_height, filter_width, input_depth, output_depth], тогда как TensorFlow.js требует [filter_height, filter_width, in_channels, out_channels]. Для плотных слоев (fc) массивы весов нужно транспонировать перед загрузкой.

Преобразование весов

function convertWeights(convnetjsParams) {
    var weights = [];
    for (var i = 0; i < convnetjsParams.length; i++) {
        var param = convnetjsParams[i];
        if (param.w) {
            // Для свёрточных слоев reshape и транспонирование
            var w = param.w; // массив Float32
            var shape = param.w_shape; // [sx, sy, in_depth, out_depth]
            // Создать tf.tensor4d из массива
            var tfTensor = tf.tensor4d(w, shape);
            weights.push(tfTensor);
        }
        if (param.b) {
            var bTensor = tf.tensor1d(param.b);
            weights.push(bTensor);
        }
    }
    return weights;
}

Создание модели TensorFlow.js

const tfLayersModel = tf.sequential();

layer_defs.forEach(layer => {
    switch(layer.type) {
        case 'conv':
            tfLayersModel.add(tf.layers.conv2d({
                inputShape: layer.out_sx && layer.out_sy ? [layer.out_sx, layer.out_sy, layer.out_depth] : undefined,
                filters: layer.filters,
                kernelSize: layer.sx,
                strides: layer.stride,
                padding: layer.pad > 0 ? 'same' : 'valid',
                activation: layer.activation
            }));
            break;
        case 'pool':
            tfLayersModel.add(tf.layers.maxPooling2d({
                poolSize: [layer.sx, layer.sx],
                strides: layer.stride
            }));
            break;
        case 'fc':
            tfLayersModel.add(tf.layers.dense({
                units: layer.num_neurons,
                activation: layer.activation
            }));
            break;
    }
});

После добавления слоев веса импортируются через setWeights().

tfLayersModel.layers.forEach((layer, i) => {
    layer.setWeights([convertedWeights[i * 2], convertedWeights[i * 2 + 1]]);
});

Сохранение TensorFlow.js модели

Модель можно сохранить локально или на сервере в формате JSON + бинарные веса:

await tfLayersModel.save('downloads://converted-model');

Формат сохраняется в виде:

  • model.json — описание структуры сети
  • weights.bin — бинарные данные весов

Этот формат полностью совместим с TensorFlow.js и позволяет использовать модель для инференса в браузере или Node.js.

Особенности и рекомендации

  • ConvNetJS не поддерживает Batch Normalization и Dropout как отдельные слои для экспортирования в TensorFlow.js, их нужно обрабатывать вручную.
  • Функции активации должны строго соответствовать TensorFlow.js, иначе возможны расхождения в поведении сети.
  • Для больших сетей желательно использовать TypedArray (Float32Array) при конвертации весов для уменьшения памяти и ускорения загрузки.
  • При переносе моделей необходимо тестировать предсказания на небольшом наборе данных, чтобы убедиться в корректности преобразования.

Экспорт из ConvNetJS в TensorFlow.js открывает возможность использовать существующие прототипы в производственных приложениях с GPU-ускорением и расширенными инструментами визуализации.

Оглавление