# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="j8Pi6JMTQn4y" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 34} outputId="aa8605d6-42e0-4e05-b41f-149849ebd8c6"
import numpy as np

import tensorflow as tf
import keras.backend as K
from keras.datasets import mnist, cifar10
from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, BatchNormalization, MaxPooling2D
from keras.layers import Flatten
from keras.optimizers import SGD, Adam, RMSprop
from keras.callbacks import LearningRateScheduler
from keras.utils import np_utils


# + id="BFsfPsqIQzJR" colab_type="code" colab={}
def round_through(x):
    '''Element-wise rounding to the closest integer with full gradient propagation.
    A trick from [<NAME>](http://stackoverflow.com/a/36480182)
    '''
    rounded = K.round(x)
    return x + K.stop_gradient(rounded - x)


def _hard_sigmoid(x):
    '''Hard sigmoid different from the more conventional form (see definition of K.hard_sigmoid).

    # Reference:
    - [BinaryNet: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1, Courbariaux et al. 2016](http://arxiv.org/abs/1602.02830}

    '''
    x = (0.5 * x) + 0.5
    return K.clip(x, 0, 1)


def binary_sigmoid(x):
    '''Binary hard sigmoid for training binarized neural network.

    # Reference:
    - [BinaryNet: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1, Courbariaux et al. 2016](http://arxiv.org/abs/1602.02830}

    '''
    return round_through(_hard_sigmoid(x))


def binary_tanh(x):
    '''Binary hard sigmoid for training binarized neural network.
     The neurons' and weights' activations binarization function
     It behaves like the sign function during forward propagation
     And like:
        hard_tanh(x) = 2 * _hard_sigmoid(x) - 1 
        clear gradient when |x| > 1 during back propagation

    # Reference:
    - [BinaryNet: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1, Courbariaux et al. 2016](http://arxiv.org/abs/1602.02830}

    '''
    return 2 * round_through(_hard_sigmoid(x)) - 1


# + id="Ro2y-LqdRAS9" colab_type="code" colab={}
from keras.layers import Dense, Conv2D
from keras import constraints
from keras import initializers


'''Binarized Dense and Convolution2D layers
References: 
"BinaryNet: Training Deep Neural Networks with Weights and Activations Constrained to +1 or -1" [http://arxiv.org/abs/1602.02830]
'''


class Clip(constraints.Constraint):

    def __call__(self, p):
        return K.clip(p, -1., 1.)


class BinaryDense(Dense):

    def __init__(self, units, **kwargs):
        super(BinaryDense, self).__init__(units, use_bias=False, 
                                          kernel_initializer=initializers.RandomUniform(-1., 1.),
                                          kernel_constraint=Clip(), **kwargs)


    def call(self, inputs):
        
        # 1. Binarize weights
        binary_kernel = binary_tanh(self.kernel)
        
        # 2. Perform matrix multiplication
        output = K.dot(inputs, binary_kernel)
        return output


class BinaryConv2D(Conv2D):

    def __init__(self, filters, kernel_size, **kwargs):
        super(BinaryConv2D, self).__init__(filters, kernel_size, use_bias=False, padding='same',
                                           kernel_initializer=initializers.RandomUniform(-1., 1.),
                                           kernel_constraint=Clip(), **kwargs)
        
    def call(self, inputs):
        
        # 1. Binarize weights
        binary_kernel = binary_tanh(self.kernel)
        
        if self.padding == 'same':
            # ones padding instead of zero padding (to keep the whole network real binary)
            kh = self.kernel_size // 2 if type(self.kernel_size) == int else self.kernel_size[0] // 2
            inputs = tf.pad(inputs, [[0,0], [kh, kh], [kh, kh], [0,0]], constant_values=1.)
        
        # 2. Perform convolution
        outputs = K.conv2d(
            inputs,
            binary_kernel,
            strides=self.strides,
            padding='valid',
            data_format=self.data_format,
            dilation_rate=self.dilation_rate)
        
        return outputs


# + id="y65wFEzJRG2S" colab_type="code" colab={}
epochs = 20
# learning rate schedule
lr_start = 1e-3
lr_end = 1e-4
lr_decay = (lr_end / lr_start)**(1. / epochs)

# BN
epsilon = 1e-6  # default 1e-3
momentum = 0.9  # default 0.99

# + id="1rv1YlKKRIzo" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 51} outputId="377c3d69-0700-4a80-9567-ce6b77165c25"
# the data, shuffled and split between train and test sets
dataset = 'mnist'
train = 'softmax'


if dataset == 'mnist':
    (X_train, y_train), (X_test, y_test) = mnist.load_data()
    X_train = X_train.reshape(60000, 28, 28, 1)
    X_test = X_test.reshape(10000, 28, 28, 1)
elif dataset == 'cifar10':
    (X_train, y_train), (X_test, y_test) = cifar10.load_data()
    
X_train = X_train.astype('float32') / 255.
X_test = X_test.astype('float32') / 255.

print(X_train.shape[0], 'train samples')
print(X_test.shape[0], 'test samples')

# convert class vectors to binary class matrices
classes = y_train.max() + 1

if train == 'hinge':
    Y_train = np_utils.to_categorical(y_train, classes) * 2 - 1 # -1 or 1 for hinge loss
    Y_test = np_utils.to_categorical(y_test, classes) * 2 - 1
elif train == 'softmax':
    Y_train = Y_train = np_utils.to_categorical(y_train, classes)
    Y_test = np_utils.to_categorical(y_test, classes)

# + id="gOD4WemORLjm" colab_type="code" colab={}
model = Sequential()


# conv1
model.add(BinaryConv2D(32, (3,3), name='conv1', input_shape=X_train.shape[1:]))
model.add(BatchNormalization(epsilon=epsilon, momentum=momentum, name='bn1'))
model.add(Activation(binary_tanh, name='act1'))
# conv2
model.add(BinaryConv2D(32, (3,3), name='conv2'))
model.add(BatchNormalization(epsilon=epsilon, momentum=momentum, name='bn2'))
model.add(Activation(binary_tanh, name='act2'))

model.add(MaxPooling2D(name='pool2'))

# conv3
model.add(BinaryConv2D(64, (3,3), name='conv3'))
model.add(BatchNormalization(epsilon=epsilon, momentum=momentum, name='bn3'))
model.add(Activation(binary_tanh, name='act3'))
# conv4
model.add(BinaryConv2D(64, (3,3), name='conv4'))
model.add(BatchNormalization(epsilon=epsilon, momentum=momentum, name='bn4'))
model.add(Activation(binary_tanh, name='act4'))

model.add(MaxPooling2D(name='pool4'))

model.add(Flatten())
# dense1
model.add(BinaryDense(128, name='dense5'))
model.add(BatchNormalization(epsilon=epsilon, momentum=momentum, name='bn5'))
model.add(Activation(binary_tanh, name='act5'))
# dense2
model.add(BinaryDense(classes, name='dense6'))
model.add(BatchNormalization(epsilon=epsilon, momentum=momentum, name='bn6'))

if train =='softmax':
    model.add(Activation('softmax'))

# + id="jg_on0asRPqC" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 850} outputId="493ca386-637c-4866-edc4-d971e6da8113"
opt = Adam(lr=lr_start)

if train == 'softmax':
    loss = 'categorical_crossentropy'
elif train == 'hinge':
    loss = 'squared_hinge'

model.compile(loss=loss, optimizer=opt, metrics=['acc'])
model.summary()

lr_scheduler = LearningRateScheduler(lambda e: lr_start * lr_decay ** e)

# + id="kIWxmA2kRTDp" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 714} outputId="8a30c075-1700-41b8-8a5d-b86ffe3ec34a"
history = model.fit(X_train, Y_train,
                    validation_data=(X_test, Y_test),
                    batch_size=256, epochs=epochs,
                    callbacks=[lr_scheduler])

# + id="b2fAO6dmvrQy" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 364} outputId="3724f4f9-1515-4758-fc27-7d2b90f6a0b9"
import matplotlib.pyplot as plt

plt.plot(history.history['acc'], label='train')
plt.plot(history.history['val_acc'], label='test')
plt.grid()
plt.legend()
plt.plot()

StarcoderdataJupyterScriptsDedupFiltered