This example demonstrates how to write custom layers for Keras.

We build a custom activation layer called 'Antirectifier',which modifies the shape of the tensor that passes through it.We need to specify two methods: compute_output_shape and call.

Note that the same result can also be achieved via a Lambda layer.

Because our custom layer is written with primitives from the Kerasbackend (K), our code can run both on TensorFlow and Theano.

  1. from __future__ import print_function
  2. import keras
  3. from keras.models import Sequential
  4. from keras import layers
  5. from keras.datasets import mnist
  6. from keras import backend as K
  9. class Antirectifier(layers.Layer):
  10.     '''This is the combination of a sample-wise
  11.     L2 normalization with the concatenation of the
  12.     positive part of the input with the negative part
  13.     of the input. The result is a tensor of samples that are
  14.     twice as large as the input samples.
  16.     It can be used in place of a ReLU.
  18.     # Input shape
  19.         2D tensor of shape (samples, n)
  21.     # Output shape
  22.         2D tensor of shape (samples, 2*n)
  24.     # Theoretical justification
  25.         When applying ReLU, assuming that the distribution
  26.         of the previous output is approximately centered around 0.,
  27.         you are discarding half of your input. This is inefficient.
  29.         Antirectifier allows to return all-positive outputs like ReLU,
  30.         without discarding any data.
  32.         Tests on MNIST show that Antirectifier allows to train networks
  33.         with twice less parameters yet with comparable
  34.         classification accuracy as an equivalent ReLU-based network.
  35.     '''
  37.     def compute_output_shape(self, input_shape):
  38.         shape = list(input_shape)
  39.         assert len(shape) == 2  # only valid for 2D tensors
  40.         shape[-1] *= 2
  41.         return tuple(shape)
  43.     def call(self, inputs):
  44.         inputs -= K.mean(inputs, axis=1, keepdims=True)
  45.         inputs = K.l2_normalize(inputs, axis=1)
  46.         pos = K.relu(inputs)
  47.         neg = K.relu(-inputs)
  48.         return K.concatenate([pos, neg], axis=1)
  50. # global parameters
  51. batch_size = 128
  52. num_classes = 10
  53. epochs = 40
  55. # the data, split between train and test sets
  56. (x_train, y_train), (x_test, y_test) = mnist.load_data()
  58. x_train = x_train.reshape(60000, 784)
  59. x_test = x_test.reshape(10000, 784)
  60. x_train = x_train.astype('float32')
  61. x_test = x_test.astype('float32')
  62. x_train /= 255
  63. x_test /= 255
  64. print(x_train.shape[0], 'train samples')
  65. print(x_test.shape[0], 'test samples')
  67. # convert class vectors to binary class matrices
  68. y_train = keras.utils.to_categorical(y_train, num_classes)
  69. y_test = keras.utils.to_categorical(y_test, num_classes)
  71. # build the model
  72. model = Sequential()
  73. model.add(layers.Dense(256, input_shape=(784,)))
  74. model.add(Antirectifier())
  75. model.add(layers.Dropout(0.1))
  76. model.add(layers.Dense(256))
  77. model.add(Antirectifier())
  78. model.add(layers.Dropout(0.1))
  79. model.add(layers.Dense(num_classes))
  80. model.add(layers.Activation('softmax'))
  82. # compile the model
  83. model.compile(loss='categorical_crossentropy',
  84.               optimizer='rmsprop',
  85.               metrics=['accuracy'])
  87. # train the model
  88. model.fit(x_train, y_train,
  89.           batch_size=batch_size,
  90.           epochs=epochs,
  91.           verbose=1,
  92.           validation_data=(x_test, y_test))
  94. # next, compare with an equivalent network
  95. # with2x bigger Dense layers and ReLU