title simple custom layer example antirectifier fchollethttps twitter comfchollet date created 20160106 last modified 20200420 description demonstration of custom layer creation accelerator gpu introduction this example shows how to create custom layers using the antirectifier layer originally proposed as a keras example script in january 2016 an alternative to relu instead of zeroingout the negative part of the input it splits the negative and positive parts and returns the concatenation of the absolute value of both this avoids loss of information at the cost of an increase in dimensionality to fix the dimensionality increase we linearly combine the features back to a space of the original size setup the antirectifier layer implement getconfig to enable serialization this is optional let s testdrive it on mnist training parameters the data split between train and test sets build the model compile the model train the model test the model introduction this example shows how to create custom layers using the antirectifier layer originally proposed as a keras example script in january 2016 an alternative to relu instead of zeroing out the negative part of the input it splits the negative and positive parts and returns the concatenation of the absolute value of both this avoids loss of information at the cost of an increase in dimensionality to fix the dimensionality increase we linearly combine the features back to a space of the original size setup the antirectifier layer implement get_config to enable serialization this is optional let s test drive it on mnist training parameters the data split between train and test sets build the model compile the model train the model test the model
import tensorflow as tf
import keras
from keras import layers




class Antirectifier(layers.Layer):
    def __init__(self, initializer="he_normal", **kwargs):
        super().__init__(**kwargs)
        self.initializer = keras.initializers.get(initializer)

    def build(self, input_shape):
        output_dim = input_shape[-1]
        self.kernel = self.add_weight(
            shape=(output_dim * 2, output_dim),
            initializer=self.initializer,
            name="kernel",
            trainable=True,
        )

    def call(self, inputs):
        inputs -= tf.reduce_mean(inputs, axis=-1, keepdims=True)
        pos = tf.nn.relu(inputs)
        neg = tf.nn.relu(-inputs)
        concatenated = tf.concat([pos, neg], axis=-1)
        mixed = tf.matmul(concatenated, self.kernel)
        return mixed

    def get_config(self):
        
        base_config = super().get_config()
        config = {"initializer": keras.initializers.serialize(self.initializer)}
        return dict(list(base_config.items()) + list(config.items()))





batch_size = 128
num_classes = 10
epochs = 20


(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.reshape(-1, 784)
x_test = x_test.reshape(-1, 784)
x_train = x_train.astype("float32")
x_test = x_test.astype("float32")
x_train /= 255
x_test /= 255
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")


model = keras.Sequential(
    [
        keras.Input(shape=(784,)),
        layers.Dense(256),
        Antirectifier(),
        layers.Dense(256),
        Antirectifier(),
        layers.Dropout(0.5),
        layers.Dense(10),
    ]
)


model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop(),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
)


model.fit(
    x_train,
    y_train,
    batch_size=batch_size,
    epochs=epochs,
    validation_split=0.15,
)


model.evaluate(x_test, y_test)