diff --git a/docs/modules/activation.rst b/docs/modules/activation.rst index c0c2b2223..8d7474674 100644 --- a/docs/modules/activation.rst +++ b/docs/modules/activation.rst @@ -49,7 +49,7 @@ Swish ------------ .. autofunction:: swish -Differentiable Sign +Sign --------------------- .. autofunction:: sign diff --git a/example/tutorial_binarynet_mnist_cnn.py b/example/tutorial_binarynet_mnist_cnn.py new file mode 100644 index 000000000..326c0aa31 --- /dev/null +++ b/example/tutorial_binarynet_mnist_cnn.py @@ -0,0 +1,103 @@ +#! /usr/bin/python +# -*- coding: utf-8 -*- + +import time +import tensorflow as tf +import tensorlayer as tl + +X_train, y_train, X_val, y_val, X_test, y_test = \ + tl.files.load_mnist_dataset(shape=(-1, 28, 28, 1)) + +sess = tf.InteractiveSession() + +batch_size = 128 + +x = tf.placeholder(tf.float32, shape=[batch_size, 28, 28, 1]) +y_ = tf.placeholder(tf.int64, shape=[batch_size]) + + +def model(x, is_train=True, reuse=False): + with tf.variable_scope("binarynet", reuse=reuse): + net = tl.layers.InputLayer(x, name='input') + net = tl.layers.BinaryConv2d(net, 32, (5, 5), (1, 1), padding='SAME', name='bcnn1') + net = tl.layers.MaxPool2d(net, (2, 2), (2, 2), padding='SAME', name='pool1') + + net = tl.layers.BatchNormLayer(net, is_train=is_train, name='bn') + net = tl.layers.SignLayer(net, name='sign2') + net = tl.layers.BinaryConv2d(net, 64, (5, 5), (1, 1), padding='SAME', name='bcnn2') + net = tl.layers.MaxPool2d(net, (2, 2), (2, 2), padding='SAME', name='pool2') + + net = tl.layers.SignLayer(net, name='sign2') + net = tl.layers.FlattenLayer(net, name='flatten') + net = tl.layers.DropoutLayer(net, 0.5, True, is_train, name='drop1') + # net = tl.layers.DenseLayer(net, 256, act=tf.nn.relu, name='dense') + net = tl.layers.BinaryDenseLayer(net, 256, name='dense') + net = tl.layers.DropoutLayer(net, 0.5, True, is_train, name='drop2') + # net = tl.layers.DenseLayer(net, 10, act=tf.identity, name='output') + net = tl.layers.BinaryDenseLayer(net, 10, name='bout') + # net = tl.layers.ScaleLayer(net, name='scale') + return net + + +# define inferences +net_train = model(x, is_train=True, reuse=False) +net_test = model(x, is_train=False, reuse=True) + +# cost for training +y = net_train.outputs +cost = tl.cost.cross_entropy(y, y_, name='xentropy') + +# cost and accuracy for evalution +y2 = net_test.outputs +cost_test = tl.cost.cross_entropy(y2, y_, name='xentropy2') +correct_prediction = tf.equal(tf.argmax(y2, 1), y_) +acc = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) + +# define the optimizer +train_params = tl.layers.get_variables_with_name('binarynet', True, True) +train_op = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(cost, var_list=train_params) + +# initialize all variables in the session +tl.layers.initialize_global_variables(sess) + +net_train.print_params() +net_train.print_layers() + +n_epoch = 200 +print_freq = 5 + +# print(sess.run(net_test.all_params)) # print real value of parameters + +for epoch in range(n_epoch): + start_time = time.time() + for X_train_a, y_train_a in tl.iterate.minibatches(X_train, y_train, batch_size, shuffle=True): + sess.run(train_op, feed_dict={x: X_train_a, y_: y_train_a}) + + if epoch + 1 == 1 or (epoch + 1) % print_freq == 0: + print("Epoch %d of %d took %fs" % (epoch + 1, n_epoch, time.time() - start_time)) + train_loss, train_acc, n_batch = 0, 0, 0 + for X_train_a, y_train_a in tl.iterate.minibatches(X_train, y_train, batch_size, shuffle=True): + err, ac = sess.run([cost_test, acc], feed_dict={x: X_train_a, y_: y_train_a}) + train_loss += err + train_acc += ac + n_batch += 1 + print(" train loss: %f" % (train_loss / n_batch)) + print(" train acc: %f" % (train_acc / n_batch)) + val_loss, val_acc, n_batch = 0, 0, 0 + for X_val_a, y_val_a in tl.iterate.minibatches(X_val, y_val, batch_size, shuffle=True): + err, ac = sess.run([cost_test, acc], feed_dict={x: X_val_a, y_: y_val_a}) + val_loss += err + val_acc += ac + n_batch += 1 + print(" val loss: %f" % (val_loss / n_batch)) + print(" val acc: %f" % (val_acc / n_batch)) + +print('Evaluation') +test_loss, test_acc, n_batch = 0, 0, 0 +for X_test_a, y_test_a in tl.iterate.minibatches(X_test, y_test, batch_size, shuffle=True): + err, ac = sess.run([cost_test, acc], feed_dict={x: X_test_a, y_: y_test_a}) + test_loss += err + test_acc += ac + n_batch += 1 +print(" test loss: %f" % (test_loss / n_batch)) +print(" test acc: %f" % (test_acc / n_batch)) diff --git a/tensorlayer/activation.py b/tensorlayer/activation.py index 70818b875..9a0c12896 100644 --- a/tensorlayer/activation.py +++ b/tensorlayer/activation.py @@ -123,7 +123,9 @@ def _sign_grad(unused_op, grad): def sign(x): # https://github.com/AngusG/tensorflow-xnor-bnn/blob/master/models/binary_net.py#L36 - """Differentiable sign function by clipping linear gradient into [-1, 1], usually be used for quantizing value in binary network, see `tf.sign `__. + """Sign function. + + Clip and binarize tensor using the straight through estimator (STE) for the gradient, usually be used for quantizing values in `Binarized Neural Networks `__. Parameters ---------- @@ -141,7 +143,7 @@ def sign(x): # https://github.com/AngusG/tensorflow-xnor-bnn/blob/master/models """ with tf.get_default_graph().gradient_override_map({"sign": "QuantizeGrad"}): - return tf.sign(x, name='tl_sign') + return tf.sign(x, name='sign') # if tf.__version__ > "1.7": diff --git a/tensorlayer/layers/__init__.py b/tensorlayer/layers/__init__.py index cad53aab8..66b0d9736 100644 --- a/tensorlayer/layers/__init__.py +++ b/tensorlayer/layers/__init__.py @@ -9,6 +9,7 @@ from .core import * from .convolution import * +from .binary import * from .super_resolution import * from .normalization import * from .spatial_transformer import * diff --git a/tensorlayer/layers/binary.py b/tensorlayer/layers/binary.py new file mode 100644 index 000000000..9d0578ee3 --- /dev/null +++ b/tensorlayer/layers/binary.py @@ -0,0 +1,275 @@ +# -*- coding: utf-8 -*- +from .core import * +from .. import _logging as logging +import tensorflow as tf + +__all__ = [ + 'BinaryDenseLayer', + 'SignLayer', + 'ScaleLayer', + 'BinaryConv2d', +] + + +@tf.RegisterGradient("TL_Sign_QuantizeGrad") +def _quantize_grad(op, grad): + """Clip and binarize tensor using the straight through estimator (STE) for the gradient. """ + return tf.clip_by_value(tf.identity(grad), -1, 1) + + +def quantize(x): + # ref: https://github.com/AngusG/tensorflow-xnor-bnn/blob/master/models/binary_net.py#L70 + # https://github.com/itayhubara/BinaryNet.tf/blob/master/nnUtils.py + with tf.get_default_graph().gradient_override_map({"Sign": "TL_Sign_QuantizeGrad"}): + return tf.sign(x) + + +class BinaryDenseLayer(Layer): + """The :class:`BinaryDenseLayer` class is a binary fully connected layer, which weights are either -1 or 1 while inferencing. + + Note that, the bias vector would not be binarized. + + Parameters + ---------- + layer : :class:`Layer` + Previous layer. + n_units : int + The number of units of this layer. + act : activation function + The activation function of this layer, usually set to ``tf.act.sign`` or apply :class:`SignLayer` after :class:`BatchNormLayer`. + use_gemm : boolean + If True, use gemm instead of ``tf.matmul`` for inferencing. (TODO). + W_init : initializer + The initializer for the weight matrix. + b_init : initializer or None + The initializer for the bias vector. If None, skip biases. + W_init_args : dictionary + The arguments for the weight matrix initializer. + b_init_args : dictionary + The arguments for the bias vector initializer. + name : a str + A unique layer name. + + """ + + def __init__( + self, + prev_layer, + n_units=100, + act=tf.identity, + use_gemm=False, + W_init=tf.truncated_normal_initializer(stddev=0.1), + b_init=tf.constant_initializer(value=0.0), + W_init_args=None, + b_init_args=None, + name='binary_dense', + ): + if W_init_args is None: + W_init_args = {} + if b_init_args is None: + b_init_args = {} + + Layer.__init__(self, prev_layer=prev_layer, name=name) + self.inputs = prev_layer.outputs + if self.inputs.get_shape().ndims != 2: + raise Exception("The input dimension must be rank 2, please reshape or flatten it") + + if use_gemm: + raise Exception("TODO. The current version use tf.matmul for inferencing.") + + n_in = int(self.inputs.get_shape()[-1]) + self.n_units = n_units + logging.info("BinaryDenseLayer %s: %d %s" % (self.name, self.n_units, act.__name__)) + with tf.variable_scope(name): + W = tf.get_variable(name='W', shape=(n_in, n_units), initializer=W_init, dtype=LayersConfig.tf_dtype, **W_init_args) + # W = tl.act.sign(W) # dont update ... + W = quantize(W) + # W = tf.Variable(W) + # print(W) + if b_init is not None: + try: + b = tf.get_variable(name='b', shape=(n_units), initializer=b_init, dtype=LayersConfig.tf_dtype, **b_init_args) + except Exception: # If initializer is a constant, do not specify shape. + b = tf.get_variable(name='b', initializer=b_init, dtype=LayersConfig.tf_dtype, **b_init_args) + self.outputs = act(tf.matmul(self.inputs, W) + b) + # self.outputs = act(xnor_gemm(self.inputs, W) + b) # TODO + else: + self.outputs = act(tf.matmul(self.inputs, W)) + # self.outputs = act(xnor_gemm(self.inputs, W)) # TODO + + self.all_layers.append(self.outputs) + if b_init is not None: + self.all_params.extend([W, b]) + else: + self.all_params.append(W) + + +class BinaryConv2d(Layer): + """ + The :class:`BinaryConv2d` class is a 2D binary CNN layer, which weights are either -1 or 1 while inferencing. + + Note that, the bias vector would not be binarized. + + Parameters + ---------- + layer : :class:`Layer` + Previous layer. + n_filter : int + The number of filters. + filter_size : tuple of int + The filter size (height, width). + strides : tuple of int + The sliding window strides of corresponding input dimensions. + It must be in the same order as the ``shape`` parameter. + act : activation function + The activation function of this layer. + padding : str + The padding algorithm type: "SAME" or "VALID". + use_gemm : boolean + If True, use gemm instead of ``tf.matmul`` for inferencing. (TODO). + W_init : initializer + The initializer for the the weight matrix. + b_init : initializer or None + The initializer for the the bias vector. If None, skip biases. + W_init_args : dictionary + The arguments for the weight matrix initializer. + b_init_args : dictionary + The arguments for the bias vector initializer. + use_cudnn_on_gpu : bool + Default is False. + data_format : str + "NHWC" or "NCHW", default is "NHWC". + name : str + A unique layer name. + + """ + + def __init__( + self, + prev_layer, + n_filter=32, + filter_size=(3, 3), + strides=(1, 1), + act=tf.identity, + padding='SAME', + use_gemm=False, + W_init=tf.truncated_normal_initializer(stddev=0.02), + b_init=tf.constant_initializer(value=0.0), + W_init_args=None, + b_init_args=None, + use_cudnn_on_gpu=None, + data_format=None, + # act=tf.identity, + # shape=(5, 5, 1, 100), + # strides=(1, 1, 1, 1), + # padding='SAME', + # W_init=tf.truncated_normal_initializer(stddev=0.02), + # b_init=tf.constant_initializer(value=0.0), + # W_init_args=None, + # b_init_args=None, + # use_cudnn_on_gpu=None, + # data_format=None, + name='binary_cnn2d', + ): + if W_init_args is None: + W_init_args = {} + if b_init_args is None: + b_init_args = {} + + if use_gemm: + raise Exception("TODO. The current version use tf.matmul for inferencing.") + + Layer.__init__(self, prev_layer=prev_layer, name=name) + self.inputs = prev_layer.outputs + if act is None: + act = tf.identity + logging.info("BinaryConv2d %s: n_filter:%d filter_size:%s strides:%s pad:%s act:%s" % (self.name, n_filter, str(filter_size), str(strides), padding, + act.__name__)) + + if len(strides) != 2: + raise ValueError("len(strides) should be 2.") + try: + pre_channel = int(prev_layer.outputs.get_shape()[-1]) + except Exception: # if pre_channel is ?, it happens when using Spatial Transformer Net + pre_channel = 1 + logging.info("[warnings] unknow input channels, set to 1") + shape = (filter_size[0], filter_size[1], pre_channel, n_filter) + strides = (1, strides[0], strides[1], 1) + with tf.variable_scope(name): + W = tf.get_variable(name='W_conv2d', shape=shape, initializer=W_init, dtype=LayersConfig.tf_dtype, **W_init_args) + W = quantize(W) + if b_init: + b = tf.get_variable(name='b_conv2d', shape=(shape[-1]), initializer=b_init, dtype=LayersConfig.tf_dtype, **b_init_args) + self.outputs = act( + tf.nn.conv2d(self.inputs, W, strides=strides, padding=padding, use_cudnn_on_gpu=use_cudnn_on_gpu, data_format=data_format) + b) + else: + self.outputs = act(tf.nn.conv2d(self.inputs, W, strides=strides, padding=padding, use_cudnn_on_gpu=use_cudnn_on_gpu, data_format=data_format)) + + self.all_layers.append(self.outputs) + if b_init: + self.all_params.extend([W, b]) + else: + self.all_params.append(W) + + +class SignLayer(Layer): + """The :class:`SignLayer` class is for quantizing the layer outputs to -1 or 1 while inferencing. + + Parameters + ---------- + layer : :class:`Layer` + Previous layer. + name : a str + A unique layer name. + + """ + + def __init__( + self, + prev_layer, + name='sign', + ): + + Layer.__init__(self, prev_layer=prev_layer, name=name) + self.inputs = prev_layer.outputs + + logging.info("SignLayer %s" % (self.name)) + with tf.variable_scope(name): + # self.outputs = tl.act.sign(self.inputs) + self.outputs = quantize(self.inputs) + + self.all_layers.append(self.outputs) + + +class ScaleLayer(Layer): + """The :class:`AddScaleLayer` class is for multipling a trainble scale value to the layer outputs. Usually be used on the output of binary net. + + Parameters + ---------- + layer : :class:`Layer` + Previous layer. + init_scale : float + The initial value for the scale factor. + name : a str + A unique layer name. + + """ + + def __init__( + self, + prev_layer, + init_scale=0.05, + name='scale', + ): + + Layer.__init__(self, prev_layer=prev_layer, name=name) + self.inputs = prev_layer.outputs + + logging.info("ScaleLayer %s: init_scale: %f" % (self.name, init_scale)) + with tf.variable_scope(name): + # scale = tf.get_variable(name='scale_factor', init, trainable=True, ) + scale = tf.get_variable("scale", shape=[1], initializer=tf.constant_initializer(value=init_scale)) + self.outputs = self.inputs * scale + + self.all_layers.append(self.outputs) + self.all_params.append(scale)