Написание пользовательского слоя

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

Базовая структура пользовательского слоя

Для создания собственного слоя необходимо определить класс, наследующий от базового класса Layer. Минимальные элементы, которые требуется реализовать:

  • Конструктор — инициализация параметров слоя.
  • Метод call(inputs) — основная логика вычислений. Принимает входные данные в виде массива Tensor и возвращает результат работы слоя.
  • Метод computeOutputShape(inputShape) — определяет форму выходного тензора на основе формы входного.

Пример базового шаблона:

class CustomLayer extends KerasJS.Layer {
  constructor(config) {
    super(config);
    this.units = config.units || 1;
    this.activation = config.activation || null;
  }

  computeOutputShape(inputShape) {
    return [inputShape[0], this.units];
  }

  call(inputs) {
    const input = inputs[0]; // inputs всегда массив
    // Простейшая линейная операция: output = input * weight + bias
    let output = input.map(x => x * this.units + 1); // пример
    if (this.activation) {
      output = this.activation(output);
    }
    return output;
  }
}

Ключевые моменты:

  • inputs всегда массив, даже если слой имеет один вход.
  • Любые параметры слоя (веса, смещения, коэффициенты) должны храниться как свойства класса.
  • Для активации можно использовать встроенные функции Keras.js (KerasJS.activations.relu, KerasJS.activations.sigmoid и др.).

Определение и управление весами

Пользовательские слои часто требуют своих параметров (весов и смещений). В Keras.js веса хранятся в объекте this.weights. Чтобы слой корректно работал, необходимо:

  1. Определить массивы для весов.
  2. Инициализировать их с помощью случайного распределения или констант.
  3. Подключить веса к графу модели.

Пример инициализации весов:

class DenseCustom extends KerasJS.Layer {
  constructor(config) {
    super(config);
    this.units = config.units;
    this.inputDim = config.inputDim;
    // Инициализация весов
    this.weights = [
      { name: 'W', data: new Float32Array(this.inputDim * this.units).fill(0.01) },
      { name: 'b', data: new Float32Array(this.units).fill(0) }
    ];
  }

  call(inputs) {
    const input = inputs[0];
    const W = this.weights[0].data;
    const b = this.weights[1].data;
    const output = new Float32Array(this.units);
    for (let i = 0; i < this.units; i++) {
      let sum = b[i];
      for (let j = 0; j < this.inputDim; j++) {
        sum += input[j] * W[j * this.units + i];
      }
      output[i] = sum;
    }
    return output;
  }

  computeOutputShape(inputShape) {
    return [inputShape[0], this.units];
  }
}

Особенности:

  • Доступ к весам осуществляется через объект this.weights.
  • Для многомерных операций важно корректно вычислять индексы при линейной алгебре.
  • Выходной массив необходимо создавать заранее, чтобы избежать лишнего выделения памяти внутри цикла.

Использование пользовательского слоя в модели

После создания слоя его можно интегрировать в модель так же, как стандартный слой Keras.js. Модель строится с помощью объекта Sequential или Model:

const model = new KerasJS.Sequential();
model.add(new DenseCustom({ units: 4, inputDim: 3 }));
model.add(new KerasJS.layers.Activation({ activation: 'relu' }));

Особенности интеграции:

  • Параметр inputDim обязателен для первого слоя, чтобы модель знала форму входа.
  • Последующие слои могут автоматически определять форму входа на основе предыдущего слоя.
  • Пользовательский слой может содержать несколько входов и выходов, если call возвращает массив тензоров.

Обработка сложных операций

Для нестандартных функций активации или комбинаций слоев можно реализовать логику напрямую в методе call. Например, слой с умножением матриц и применением нестандартной функции:

call(inputs) {
  const x = inputs[0];
  const y = inputs[1];
  const output = new Float32Array(x.length);
  for (let i = 0; i < x.length; i++) {
    output[i] = Math.tanh(x[i] * y[i] + 0.5);
  }
  return output;
}

Рекомендации:

  • Использовать встроенные функции Keras.js для активаций и операций с тензорами, если возможно, для ускорения через WebGL.
  • Для операций с большими массивами применять векторизацию и TypedArray для повышения производительности.

Поддержка нескольких входов и выходов

Пользовательский слой может принимать несколько входов и возвращать несколько выходов. В этом случае:

  • inputs всегда массив: каждый элемент — отдельный вход.
  • Метод call должен возвращать массив соответствующей длины.
  • Веса могут быть общими для всех входов или отдельными для каждого.

Пример слоя с двумя входами:

call(inputs) {
  const [input1, input2] = inputs;
  const output = input1.map((v, i) => v + input2[i]);
  return [output];
}

Сериализация и сохранение пользовательского слоя

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

  • Класс слоя доступен в контексте исполнения.
  • Формат весов соответствует структуре, ожидаемой слоем.

Для сериализации можно реализовать методы getConfig и fromConfig:

getConfig() {
  return { units: this.units, activation: this.activation };
}

static fromConfig(config) {
  return new CustomLayer(config);
}

Это позволяет интегрировать пользовательские слои в JSON модели и загружать их без ошибок.

Оптимизация производительности

  • Использовать Float32Array для всех числовых данных.
  • Минимизировать создание временных объектов внутри циклов.
  • Для матричных операций применять библиотеки линейной алгебры с поддержкой WebGL (например, встроенные функции Keras.js).
  • Выносить константные вычисления из метода call для экономии ресурсов.

Пользовательские слои в Keras.js дают полный контроль над архитектурой модели и вычислительными процессами, позволяя создавать уникальные решения прямо в браузере без серверной части. Их грамотная реализация требует аккуратного управления весами, формами тензоров и оптимизацией под WebGL.

Оглавление