How to create custom Activation functions in Keras / TensorFlow?

Question

I'm using keras and I wanted to add my own activation function myf to tensorflow backend. how to define the new function and make it operational. so instead of the line of code:

model.add(layers.Conv2D(64, (3, 3), activation='relu')) 

I'll write

model.add(layers.Conv2D(64, (3, 3), activation='myf')). 

Answer 1

First you need to define a function using backend functions. As an example, here is how I implemented the swish activation function:

from keras import backend as K

def swish(x, beta=1.0):
    return x * K.sigmoid(beta * x)

This allows you to add the activation function to your model like this:

model.add(Conv2D(64, (3, 3)))
model.add(Activation(swish))

If you want to use a string as an alias for your custom function you will have to register the custom object with Keras. It can be done like this:

from keras.utils.generic_utils import get_custom_objects

get_custom_objects().update({'swish': Activation(swish)})

This allows you to add the activation directly to layer by name:

model.add(Conv2D(64, (3, 3), activation='swish'))

For more advanced activation functions, with trainable parameters and such, it is best to implement them as a Keras Layer. Here the swish function is used in a layer, allowing beta to be learned while training:

from keras.layers import Layer

class Swish(Layer):

    def __init__(self, beta=1.0, trainable=False, **kwargs):
        super(Swish, self).__init__(**kwargs)
        self.supports_masking = True
        self.beta = beta
        self.trainable = trainable

    def build(self, input_shape):
        self.beta_factor = K.variable(self.beta,
                                      dtype=K.floatx(),
                                      name='beta_factor')
        if self.trainable:
            self._trainable_weights.append(self.beta_factor)

        super(Swish, self).build(input_shape)

    def call(self, inputs, mask=None):
        return swish(inputs, self.beta_factor)

    def get_config(self):
        config = {'beta': self.get_weights()[0] if self.trainable else self.beta,
                  'trainable': self.trainable}
        base_config = super(Swish, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

You then would add the activation function the same as any other layer:

model.add(Conv2D(64, (3, 3)))
model.add(Swish(beta=1.0, trainable=True))

Answer 2

EDIT:
For the latest versions of TensorFlow I would suggest to write all th steps using tensorflow ops (not just from keras.backend).
Write a function if just tf.ops and then put a @tf.function decorator on top of it to do jit compilation and make it faster.

---

The trick is to use Keras' backend funcions:

from keras import backend as K
def my_function(x):
    x = K.some_function(x)
    return x

where "some_function" is what you need. Can you elaborate on that?

And then you simply call it with:

model.add(Dense(10, activation = my_function))

What activation are you trying to implement?

Answer 3

• the easy way:

from keras.layers.core import Activation
from keras.models import Sequential
import keras.backend as K
def myCustomActivation(x):
    return ...
model = Sequential()
model.add(Dense(120))
model.add(Activation(myCustomActivation))
model.add(Dense(30, activation= myCustomActivation))
...

for all operations performed in your myCustomActivation use the Keras backend to get operations executed on the DAG.

or

• the "less easy" way:

an activation is just a trivial layer, thus you can then define your custom activation as a custom layer following instructions here.

For all the operations written in the call method use the keras backend (again to get everything done on the DAG).