Привет, Хабр!

Keras предоставляет мощные инструменты для создания сложных нейронных сетей. Однако иногда стандартного набора слоев недостаточно для решения некоторых задач. В таких случаях на помощь приходят кастомные слои.

Кастомные слои позволяют адаптировать архитектуру модели под особенности данных, улучшая тем самым производительность и точность моделек.

Создание кастомных слоев

Каждый кастомный слой начинается с определения нового класса, наследующего от tf.keras.layers.Layer. В __init__ происходит инициализация слоя, где можно задать параметры, необходимые для работы слоя:

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, units=32, activation=None, **kwargs):
        super(CustomLayer, self).__init__(**kwargs)
        self.units = units
        self.activation = tf.keras.activations.get(activation)

Тут units определяет количество нейронов, а activation указывает функцию активации. super(CustomLayer, self).__init__(**kwargs) вызывает конструктор базового класса Layer.

Метод build вызывается Keras при первом использовании слоя. Его юзают для создания параметров слоя, которые зависят от размера входных данных:

    def build(self, input_shape):
        self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
                                      initializer='glorot_uniform',
                                      trainable=True)
        self.bias = self.add_weight(shape=(self.units,),
                                    initializer='zeros',
                                    trainable=True)
        super(CustomLayer, self).build(input_shape)

В методе создаются веса kernel и bias. Функция add_weight создает и регистрирует переменные слоя, которые будут обновляться во время тренировки.

Метод call содержит основную логику вычислений слоя. Он принимает входные данные и возвращает выходные:

    def call(self, inputs):
        output = tf.matmul(inputs, self.kernel) + self.bias
        if self.activation is not None:
            output = self.activation(output)
        return output

В этом методе выполняется умножение входных данных на веса и добавление смещения. Если определена функция активации, она применяется к выходным данным.

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

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(8,)),
    CustomLayer(units=64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

Другие полезные методы:

• add_weight: Добавляет переменную веса в слой.

• compute_output_shape: Возвращает форму выходных данных на основе формы входных.

• get_config: Возвращает конфигурацию слоя в виде словаря, что полезно для сериализации.

Примеры реализации: Dense, Convolutional и еще три типа других слоев

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

import tensorflow as tf

class CustomDenseLayer(tf.keras.layers.Layer):
    def __init__(self, units=32, activation=None):
        super(CustomDenseLayer, self).__init__()
        self.units = units
        self.activation = tf.keras.activations.get(activation)

    def build(self, input_shape):
        self.w = self.add_weight(shape=(input_shape[-1], self.units),
                                 initializer='glorot_uniform',
                                 trainable=True)
        self.b = self.add_weight(shape=(self.units,),
                                 initializer='zeros',
                                 trainable=True)

    def call(self, inputs):
        z = tf.matmul(inputs, self.w) + self.b
        if self.activation is not None:
            return self.activation(z)
        return z

# example
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(8,)),
    CustomDenseLayer(units=64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

Convolutional слои применяют свертку фильтра к входным данным, что позволяет выделять пространственные особенности:

class CustomConvLayer(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, strides=(1, 1), padding='valid', activation=None):
        super(CustomConvLayer, self).__init__()
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding
        self.activation = tf.keras.activations.get(activation)

    def build(self, input_shape):
        self.kernel = self.add_weight(shape=(*self.kernel_size, input_shape[-1], self.filters),
                                      initializer='glorot_uniform',
                                      trainable=True)

    def call(self, inputs):
        conv = tf.nn.conv2d(inputs, self.kernel, strides=self.strides, padding=self.padding.upper())
        if self.activation is not None:
            return self.activation(conv)
        return conv

# example
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(28, 28, 1)),
    CustomConvLayer(filters=32, kernel_size=(3, 3), activation='relu'),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

Recurrent слои используются для обработки последовательных данных. Один из наиболее распространнных типов рекуррентных слоев — это LSTM:

class CustomLSTMLayer(tf.keras.layers.Layer):
    def __init__(self, units):
        super(CustomLSTMLayer, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.lstm_cell = tf.keras.layers.LSTMCell(self.units)

    def call(self, inputs, states):
        return self.lstm_cell(inputs, states)

# example
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(None, 8)),
    tf.keras.layers.RNN(CustomLSTMLayer(units=64)),
    tf.keras.layers.Dense(10, activation='softmax')
])

Dropout слой используется для регуляризации модели, предотвращая переобучение путем случайного зануления некоторых нейронов во время тренировки:

class CustomDropoutLayer(tf.keras.layers.Layer):
    def __init__(self, rate):
        super(CustomDropoutLayer, self).__init__()
        self.rate = rate

    def call(self, inputs, training=None):
        return tf.nn.dropout(inputs, rate=self.rate) if training else inputs

# example
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(8,)),
    tf.keras.layers.Dense(64, activation='relu'),
    CustomDropoutLayer(rate=0.5),
    tf.keras.layers.Dense(10, activation='softmax')
])

BatchNormalization слой нормализует активации предыдущего слоя, улучшая скорость обучения модельки:

class CustomBatchNormalizationLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(CustomBatchNormalizationLayer, self).__init__()

    def build(self, input_shape):
        self.gamma = self.add_weight(shape=(input_shape[-1],),
                                     initializer='ones',
                                     trainable=True)
        self.beta = self.add_weight(shape=(input_shape[-1],),
                                    initializer='zeros',
                                    trainable=True)
        self.moving_mean = self.add_weight(shape=(input_shape[-1],),
                                           initializer='zeros',
                                           trainable=False)
        self.moving_variance = self.add_weight(shape=(input_shape[-1],),
                                               initializer='ones',
                                               trainable=False)

    def call(self, inputs, training=None):
        if training:
            mean, variance = tf.nn.moments(inputs, axes=[0])
            self.moving_mean.assign(self.moving_mean * 0.9 + mean * 0.1)
            self.moving_variance.assign(self.moving_variance * 0.9 + variance * 0.1)
        else:
            mean, variance = self.moving_mean, self.moving_variance

        return tf.nn.batch_normalization(inputs, mean, variance, self.beta, self.gamma, variance_epsilon=1e-3)

# example
model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(8,)),
    CustomBatchNormalizationLayer(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

Больше практических инструментов и кейсов коллеги из OTUS рассматривают в рамках практических онлайн-курсов. Напомню, что с полным каталогом курсов можно ознакомиться по ссылке.