activation function for binarized neural networks

Question

I am trying to implement a binarized neural network using keras and tensorflow. Here is my current code:

import tensorflow as tf
from tensorflow import keras
from keras import datasets
from keras.utils.vis_utils import plot_model

from keras.layers import Activation
from keras import backend as K
from keras.utils.generic_utils import get_custom_objects


def sign(x):
    result = K.sign(x)
    return result


get_custom_objects().update({'custom_activation': Activation(sign)})

digit_data = datasets.mnist
(train_images, train_labels), (test_images, test_labels) = digit_data.load_data()

# import matplotlib.pyplot as plt
# plt.imshow(train_images[0])
# print(train_labels[0])
# print(train_images[0])

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(784, Activation(sign, name='SpecialActivation')),
    keras.layers.Dense(10, activation=tf.nn.softmax)
])


# model = keras.Sequential([
#     keras.layers.Flatten(input_shape=(28, 28)),
#     keras.layers.Dense(128, activation=tf.nn.relu),
#     keras.layers.Dense(10, activation=tf.nn.softmax)
# ])

# plot_model(model, to_file='model_plot.png', show_shapes=True, show_layer_names=True)

model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

model.fit(train_images, train_labels, epochs=5)
test_loss, test_acc = model.evaluate(test_images, test_labels)
print(test_loss, test_acc)

The activation function I am currently using is a sign function which returns -1 if x < 0 | 0 if x=0 | 1 if x >0 the binarized neural network needs to have a bitwise operation function which returns +1 if x => 0 and -1 if x< 0 (from what i understand)

I'm not sure how to implement that kind of an activation function.

Answer 1

Create a custom function
As explained in Keras doc for custom Actiation function

Creating custom activations
You can also use a TensorFlow callable as an activation (in this case it should take a tensor and return a tensor of the same shape and type):
model.add(layers.Dense(64, activation=tf.nn.tanh))

def custom_bin(x):
    return tf.where(x >= 0, 1.0 , -1)

import tensorflow as tf
from tensorflow import keras
model = keras.Sequential()

model.add(keras.layers.Dense( 3, activation="linear",input_shape=(5,)))
model.add(keras.layers.Dense( 3, activation="linear"))
model.add(keras.layers.Dense( 4, activation=custom_bin))
model.add(keras.layers.Dense( 2, activation="linear"))

model.summary()

A toy NN to check the Activation output

A simple neural net with all weights = 1 and bias = 0. So, it will keep the output sign same as that of the input i.e.(+Ve, 0, -Ve).Then apply the respective Activation function in the output layer. Then predict using the 3 use-cases

def custom_bin(x):
    return tf.where(x >= 0, 1.0 , -1) 
   
import tensorflow as tf
from tensorflow import keras

#Case I - tf.math.sign 
initializer = tf.keras.initializers.Ones()
model = keras.Sequential()
model.add(keras.layers.Dense( 2, activation="linear",input_shape=(2,),kernel_initializer=initializer))
model.add(keras.layers.Dense( 1, activation=tf.math.sign,kernel_initializer=initializer))
print(model.predict([[-5,-5],[0,0],[5, 5]]))

#Case II - custom if-then-else 
initializer = tf.keras.initializers.Ones()
model = keras.Sequential()
model.add(keras.layers.Dense( 2, activation="linear",input_shape=(2,),kernel_initializer=initializer))
model.add(keras.layers.Dense( 1, activation=custom_bin, kernel_initializer=initializer))
print(model.predict([[-5,-5],[0,0],[5, 5]]))

Output
[[-1.] [ 0.] [ 1.]]
[[-1.] [ 1.] [ 1.]]