<jupyter_start><jupyter_text>BSD 3-Clause LicenseCopyright (c) 2020, <NAME>All rights reserved.  The GSLU trainable activation function---The following snippet introduces a trainable activation function called "Generalized Sigmoid Linear Unit" GSLU.The function is described by:$GSLU(x) = \frac{a + bx}{1 + e^{-cx}} + d$Where $a$, $b$, $c$, & $d$ describe the trainable Parameters, $e$ describes Euler's number and $x$ describes the input/drive to the GSLU.This simple function is able to represent several known activation functions, for example:Sigmoid: $a=1$, $b=0$, $c=1$, $d=0$TanH: $a=2$, $b=0$, $c=2$, $d=-1$Swish: $a=0$, $b=1$, $c=1$, $d=0$ReLU: $a=0$, $b=1$, $c \rightarrow \infty$, $d=0$The parameters of this function can be tuned using standard backpropagation. This snippet implements the trainable activation function as a simple tf.keras Layer, and shows the simplicity of training the GSLU trainable activation function by the example of a simple Time-Series Predicting Network.  The Simplified Sigmoidal Linear UnitWe also can think of a simplified version of this, where we only distinguish between the the swish and the sigmoidal function. For this, we will only require one parameter $a$.$SSLU(x) = \frac{a + (1 - a)x}{1 + e^{-x}}$Additionally, we could put the a parameter through a sigmoid function again, in order to not only have a useful range between 0 and 1, but over all real numbers:$SSLU(x) = \frac{\frac{1}{1 + e^{-a}} + (1 - \frac{1}{1 + e^{-a}}) x}{1 + e^{-x}}$This would mean that $SSLU(x)$ will approach a Sigmoid function for the interval $a = [0 → \infty]$ and a Swish function for the interval of $a = [0 → - \infty]$. This clear definition might help us to make sure the possible convergence space of our $a$ parameter is equally spaced, which in turn might help us find the more suitable activation function for any type of neural network.Here's a plot of the $SSLU$ over different values of $a$ and $x$:![](https://www6b3.wolframalpha.com/Calculate/MSP/MSP15631h33effa4c90351h00003b1a0a7fd4717ff9?MSPStoreType=image/gif&s=10)This should help you get an intuition on how the SSLU switches from sigmoid to swish, based on the value of $a$. Now, let's try the function out.First, we will have to import all libraries that we are going to use.<jupyter_code>import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds

import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>Define Trainable activation functions<jupyter_code>class GSLU(tf.keras.layers.Layer):

  def __init__(self):
    super(GSLU, self).__init__()

  def build(self, input_shape):  # Create the state of the layer (weights)
    w_init = tf.random_normal_initializer(mean=0., stddev=1.)
    self.a = tf.Variable(initial_value=w_init(shape=input_shape[1:]), trainable=True)
    self.b = tf.Variable(initial_value=w_init(shape=input_shape[1:]), trainable=True)
    self.c = tf.Variable(initial_value=w_init(shape=input_shape[1:]), trainable=True)
    self.d = tf.Variable(initial_value=w_init(shape=input_shape[1:]), trainable=True)

  def call(self, x): 
    x = (self.a + self.b * x ) * tf.keras.activations.sigmoid(self.c * x) + self.d
    return x


class SSLU(tf.keras.layers.Layer):

  def __init__(self, variable_axes=None, **kwargs):
    """Initialize a SSLU activation as a keras Layer.
    
    Parameters
    ----------
    variable_axes : None | int | list of int
        The layer axes over which the activation function's trainable parameters
        can be varied. If `None`, the same activation is trained for the entire 
        layer. If `int` or `list(int)` the parameters are allowed to vary over 
        the corresponding axes of the previous layer's output drive. E.g. if the
        previous layers output is of shape `[batch_size, 128, 64]`, passing 
        `variable_axes=[1, 2]` will lead to a parameter shape of `[1, 128, 64]` 
        (so the activation function is variable over axes 1 and 2).
    **kwargs
        Any additional keyword argument will be passed to tf.keras.layers.Layer.
    """
    super(SSLU, self, **kwargs).__init__()
    # store the axes that should have independent weights
    if isinstance(variable_axes, int):
      variable_axes = [variable_axes]
    self.var_axes = variable_axes

  def build(self, input_shape):
    """Build the layer parameters"""
    # determine the number of weights for a
    if self.var_axes == None:
      shape = [1]
    else:
      var_axes = [i if i >= 0 else len(input_shape) + i for i in self.var_axes]
      shape = [sh if idx in var_axes else 1 for idx, sh in enumerate(input_shape)]
    # initialize a weights
    #a_init = tf.random_normal_initializer(mean=0., stddev=1.)
    #self.a = tf.Variable(initial_value=a_init(shape=shape), trainable=True)
    self.a = self.add_weight(shape=shape, initializer='glorot_normal', trainable=True)

  def call(self, x):
    """Call the activation."""
    sig_a = tf.keras.activations.sigmoid(self.a)
    #sig_a = tf.clip_by_value(self.a, 0., 1.)
    return (sig_a + (1 - sig_a) * x) * tf.keras.activations.sigmoid(x)<jupyter_output><empty_output><jupyter_text>Generate some data<jupyter_code>train_data, test_data = tfds.load('horses_or_humans', split=['train','test'],
                                  as_supervised=True, shuffle_files=True)
categories = ["horse", "human"]

for x, y in train_data.take(2):
  plt.imshow(x)
  plt.show()
  print(f"Category: {categories[y]}\n\n");
def preproc(feature, label):
  """Downscale an image."""
  feature = tf.image.resize(feature, [32, 32], method="nearest")
  return feature, label

for x, y in train_data.map(preproc).batch(1).take(2):
  plt.imshow(x[0])
  plt.show()
  print(f"Category: {categories[y[0]]}\n\n");<jupyter_output><empty_output><jupyter_text>Prepare TensorFlow dataset.<jupyter_code>tf.keras.backend.clear_session()

model = tf.keras.Sequential([
                             tf.keras.layers.Normalization(axis=-1,),
                             tf.keras.layers.Conv2D(4, (4, 4), (2, 2)),
                             SSLU(-1),
                             tf.keras.layers.Conv2D(8, (4, 4), (2, 2)),
                             SSLU(-1),
                             tf.keras.layers.Conv2D(16, (4, 4), (2, 2)),
                             SSLU(-1),
                             tf.keras.layers.Flatten(),
                             tf.keras.layers.Dense(32),
                             SSLU(-1),
                             tf.keras.layers.Dense(64),
                             SSLU(-1),
                             tf.keras.layers.Dense(128),
                             SSLU(-1),
                             tf.keras.layers.Dense(1),
                             SSLU(-1)
                             ])

model.build(x.shape)
model.summary()
model.compile(optimizer="Adam", loss="BCE", metrics=["BinaryCrossentropy", "accuracy"])
for layer in model.layers:
  if "sslu" in layer.name:
    plt.hist(np.squeeze(layer.a))
    plt.title("Distribution of alpha values for layer: " + layer.name)
    plt.xlim([-3, 3])
    plt.xlabel("alpha")
    plt.ylabel("count")
    plt.show()
model.fit(train_data.map(preproc).batch(128), epochs=20);
model.evaluate(test_data.map(preproc).batch(128));<jupyter_output>2/2 [==============================] - 1s 228ms/step - loss: 1.2829 - binary_crossentropy: 1.2829 - accuracy: 0.8555<jupyter_text>Inspect how the GSLUs turned out<jupyter_code>for layer in model.layers:
  if "sslu" in layer.name:
    sigmo_sum = np.sum(layer.a > 0)
    swish_sum = np.sum(layer.a < 0)
    print(f"Layer {layer.name}: Sigmoids: {sigmo_sum}, swishs: {swish_sum}")
for layer in model.layers:
  if "sslu" in layer.name:
    plt.hist(np.squeeze(layer.a))
    plt.title("Distribution of alpha values for layer: " + layer.name)
    plt.xlim([-3, 3])
    plt.xlabel("alpha")
    plt.ylabel("count")
    plt.show()
x_space = np.linspace(-10, 10, 256)
names = []
plt.figure(figsize=[13, 8])
for layer in model.layers:
  if "sslu" in layer.name:
    names.append(layer.name)
    shape = layer.input_shape
    batch = np.stack([np.ones(shape[1:]) * x for x in x_space])
    output = layer(batch)
    avg_output = np.mean(output, axis=tuple(range(1, len(shape))))
    plt.plot(x_space, avg_output)

plt.legend(names, title="Average activations for:")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13379774201297143, 'fraction_numerical': 0.021979341820802308, 'mean_word_length': 3.2055555555555557, 'pattern_counts': {'":': 0, '<': 22, '<?xml version=': 0, '>': 23, 'https://': 1, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '60656', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2865, 'n_tokens_neox': 2707, 'n_words': 778, 'path': 'GSLU_trainable_activation_function.ipynb', 'repo_name': 'DiGyt/snippets', 'size': '8331'}