Merge pull request #1055 from Laicheng0830/quantized_layer

Update quantized layer format

Author: Laicheng0830

tensorlayer/layers/convolution/quan_conv_bn.py

@@ -1,11 +1,11 @@
-# /usr/bin/python
+#! /usr/bin/python
 # -*- coding: utf-8 -*-
 
+import numpy as np
 import tensorflow as tf
+import tensorlayer as tl
 from tensorflow.python.training import moving_averages
-
 from tensorlayer import logging
-from tensorlayer.decorators import deprecated_alias
 from tensorlayer.layers.core import Layer
 from tensorlayer.layers.utils import (quantize_active_overflow, quantize_weight_overflow)
 
@@ -22,8 +22,6 @@ class QuanConv2dWithBN(Layer):
 
     Parameters
     ----------
-    prev_layer : :class:`Layer`
-        Previous layer.
     n_filter : int
         The number of filters.
     filter_size : tuple of int
@@ -51,49 +49,33 @@ class QuanConv2dWithBN(Layer):
         The bits of this layer's parameter
     bitA : int
         The bits of the output of previous layer
-    decay : float
-        A decay factor for `ExponentialMovingAverage`.
-        Suggest to use a large value for large dataset.
-    epsilon : float
-        Eplison.
-    is_train : boolean
-        Is being used for training or inference.
-    beta_init : initializer or None
-        The initializer for initializing beta, if None, skip beta.
-        Usually you should not skip beta unless you know what happened.
-    gamma_init : initializer or None
-        The initializer for initializing gamma, if None, skip gamma.
     use_gemm : boolean
         If True, use gemm instead of ``tf.matmul`` for inferencing. (TODO).
     W_init : initializer
         The initializer for the the weight matrix.
     W_init_args : dictionary
         The arguments for the weight matrix initializer.
-    use_cudnn_on_gpu : bool
-        Default is False.
     data_format : str
         "NHWC" or "NCHW", default is "NHWC".
+    dilation_rate : tuple of int
+        Specifying the dilation rate to use for dilated convolution.
+    in_channels : int
+        The number of in channels.
     name : str
         A unique layer name.
 
     Examples
     ---------
-    >>> import tensorflow as tf
     >>> import tensorlayer as tl
-    >>> x = tf.placeholder(tf.float32, [None, 256, 256, 3])
-    >>> net = tl.layers.InputLayer(x, name='input')
-    >>> net = tl.layers.QuanConv2dWithBN(net, 64, (5, 5), (1, 1),  act=tf.nn.relu, padding='SAME', is_train=is_train, bitW=bitW, bitA=bitA, name='qcnnbn1')
-    >>> net = tl.layers.MaxPool2d(net, (3, 3), (2, 2), padding='SAME', name='pool1')
-    ...
-    >>> net = tl.layers.QuanConv2dWithBN(net, 64, (5, 5), (1, 1), padding='SAME', act=tf.nn.relu, is_train=is_train,  bitW=bitW, bitA=bitA, name='qcnnbn2')
-    >>> net = tl.layers.MaxPool2d(net, (3, 3), (2, 2), padding='SAME', name='pool2')
-    ...
+    >>> net = tl.layers.Input([50, 256, 256, 3])
+    >>> layer = tl.layers.QuanConv2dWithBN(n_filter=64, filter_size=(5,5),strides=(1,1),padding='SAME',name='qcnnbn1')
+    >>> print(layer)
+    >>> net = tl.layers.QuanConv2dWithBN(n_filter=64, filter_size=(5,5),strides=(1,1),padding='SAME',name='qcnnbn1')(net)
+    >>> print(net)
     """
 
-    @deprecated_alias(layer='prev_layer', end_support_version=1.9)  # TODO remove this line for the 1.9 release
     def __init__(
         self,
-        prev_layer,
         n_filter=32,
         filter_size=(3, 3),
         strides=(1, 1),
@@ -102,125 +84,150 @@ def __init__(
         decay=0.9,
         epsilon=1e-5,
         is_train=False,
-        gamma_init=tf.compat.v1.initializers.ones,
-        beta_init=tf.compat.v1.initializers.zeros,
+        gamma_init=tl.initializers.truncated_normal(stddev=0.02),
+        beta_init=tl.initializers.truncated_normal(stddev=0.02),
         bitW=8,
         bitA=8,
         use_gemm=False,
-        W_init=tf.compat.v1.initializers.truncated_normal(stddev=0.02),
+        W_init=tl.initializers.truncated_normal(stddev=0.02),
         W_init_args=None,
-        use_cudnn_on_gpu=None,
-        data_format=None,
+        data_format="channels_last",
+        dilation_rate=(1, 1),
+        in_channels=None,
         name='quan_cnn2d_bn',
     ):
-        super(QuanConv2dWithBN, self).__init__(prev_layer=prev_layer, act=act, W_init_args=W_init_args, name=name)
-
+        super(QuanConv2dWithBN, self).__init__(act=act, name=name)
+        self.n_filter = n_filter
+        self.filter_size = filter_size
+        self.strides = strides
+        self.padding = padding
+        self.decay = decay
+        self.epsilon = epsilon
+        self.is_train = is_train
+        self.gamma_init = gamma_init
+        self.beta_init = beta_init
+        self.bitW = bitW
+        self.bitA = bitA
+        self.use_gemm = use_gemm
+        self.W_init = W_init
+        self.W_init_args = W_init_args
+        self.data_format = data_format
+        self.dilation_rate = dilation_rate
+        self.in_channels = in_channels
         logging.info(
             "QuanConv2dWithBN %s: n_filter: %d filter_size: %s strides: %s pad: %s act: %s " % (
                 self.name, n_filter, filter_size, str(strides), padding,
                 self.act.__name__ if self.act is not None else 'No Activation'
             )
         )
 
-        x = self.inputs
-        self.inputs = quantize_active_overflow(self.inputs, bitA)  # Do not remove
+        if self.in_channels:
+            self.build(None)
+            self._built = True
 
         if use_gemm:
             raise Exception("TODO. The current version use tf.matmul for inferencing.")
 
         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.warning("[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.compat.v1.variable_scope(name):
-            W = tf.compat.v1.get_variable(
-                name='W_conv2d', shape=shape, initializer=W_init, dtype=LayersConfig.tf_dtype, **self.W_init_args
-            )
-
-            conv = tf.nn.conv2d(
-                x, W, strides=strides, padding=padding, use_cudnn_on_gpu=use_cudnn_on_gpu, data_format=data_format
-            )
-
-            para_bn_shape = conv.get_shape()[-1:]
-
-            if gamma_init:
-                scale_para = tf.compat.v1.get_variable(
-                    name='scale_para', shape=para_bn_shape, initializer=gamma_init, dtype=LayersConfig.tf_dtype,
-                    trainable=is_train
-                )
-            else:
-                scale_para = None
-
-            if beta_init:
-                offset_para = tf.compat.v1.get_variable(
-                    name='offset_para', shape=para_bn_shape, initializer=beta_init, dtype=LayersConfig.tf_dtype,
-                    trainable=is_train
-                )
-            else:
-                offset_para = None
-
-            moving_mean = tf.compat.v1.get_variable(
-                'moving_mean', para_bn_shape, initializer=tf.compat.v1.initializers.constant(1.),
-                dtype=LayersConfig.tf_dtype, trainable=False
+    def __repr__(self):
+        actstr = self.act.__name__ if self.act is not None else 'No Activation'
+        s = (
+            '{classname}(in_channels={in_channels}, out_channels={n_filter}, kernel_size={filter_size}'
+            ', strides={strides}, padding={padding}' + actstr
+        )
+        if self.dilation_rate != (1, ) * len(self.dilation_rate):
+            s += ', dilation={dilation_rate}'
+        if self.name is not None:
+            s += ', name=\'{name}\''
+        s += ')'
+        return s.format(classname=self.__class__.__name__, **self.__dict__)
+
+    def build(self, inputs_shape):
+        if self.data_format == 'channels_last':
+            self.data_format = 'NHWC'
+            if self.in_channels is None:
+                self.in_channels = inputs_shape[-1]
+            self._strides = [1, self.strides[0], self.strides[1], 1]
+            self._dilation_rate = [1, self.dilation_rate[0], self.dilation_rate[1], 1]
+        elif self.data_format == 'channels_first':
+            self.data_format = 'NCHW'
+            if self.in_channels is None:
+                self.in_channels = inputs_shape[1]
+            self._strides = [1, 1, self.strides[0], self.strides[1]]
+            self._dilation_rate = [1, 1, self.dilation_rate[0], self.dilation_rate[1]]
+        else:
+            raise Exception("data_format should be either channels_last or channels_first")
+
+        self.filter_shape = (self.filter_size[0], self.filter_size[1], self.in_channels, self.n_filter)
+        self.W = self._get_weights("filters", shape=self.filter_shape, init=self.W_init)
+
+        para_bn_shape = (self.n_filter, )
+        if self.gamma_init:
+            self.scale_para = self._get_weights(
+                "scale_para", shape=para_bn_shape, init=self.gamma_init, trainable=self.is_train
             )
+        else:
+            self.scale_para = None
 
-            moving_variance = tf.compat.v1.get_variable(
-                'moving_variance',
-                para_bn_shape,
-                initializer=tf.compat.v1.initializers.constant(1.),
-                dtype=LayersConfig.tf_dtype,
-                trainable=False,
+        if self.beta_init:
+            self.offset_para = self._get_weights(
+                "offset_para", shape=para_bn_shape, init=self.beta_init, trainable=self.is_train
             )
+        else:
+            self.offset_para = None
 
-            mean, variance = tf.nn.moments(x=conv, axes=list(range(len(conv.get_shape()) - 1)))
-
-            update_moving_mean = moving_averages.assign_moving_average(
-                moving_mean, mean, decay, zero_debias=False
-            )  # if zero_debias=True, has bias
-
-            update_moving_variance = moving_averages.assign_moving_average(
-                moving_variance, variance, decay, zero_debias=False
-            )  # if zero_debias=True, has bias
+        self.moving_mean = self._get_weights(
+            "moving_mean", shape=para_bn_shape, init=tl.initializers.constant(1.0), trainable=False
+        )
+        self.moving_variance = self._get_weights(
+            "moving_variance", shape=para_bn_shape, init=tl.initializers.constant(1.0), trainable=False
+        )
 
-            def mean_var_with_update():
-                with tf.control_dependencies([update_moving_mean, update_moving_variance]):
-                    return tf.identity(mean), tf.identity(variance)
+    def forward(self, inputs):
+        x = inputs
+        inputs = quantize_active_overflow(inputs, self.bitA)  # Do not remove
+        outputs = tf.nn.conv2d(
+            input=x, filters=self.W, strides=self._strides, padding=self.padding, data_format=self.data_format,
+            dilations=self._dilation_rate, name=self.name
+        )
 
-            if is_train:
-                mean, var = mean_var_with_update()
-            else:
-                mean, var = moving_mean, moving_variance
+        mean, variance = tf.nn.moments(outputs, axes=list(range(len(outputs.get_shape()) - 1)))
 
-            w_fold = _w_fold(W, scale_para, var, epsilon)
-            bias_fold = _bias_fold(offset_para, scale_para, mean, var, epsilon)
+        update_moving_mean = moving_averages.assign_moving_average(
+            self.moving_mean, mean, self.decay, zero_debias=False
+        )  # if zero_debias=True, has bias
+        update_moving_variance = moving_averages.assign_moving_average(
+            self.moving_variance, mean, self.decay, zero_debias=False
+        )  # if zero_debias=True, has bias
 
-            W = quantize_weight_overflow(w_fold, bitW)
+        if self.is_train:
+            mean, var = self.mean_var_with_update(update_moving_mean, update_moving_variance, mean, variance)
+        else:
+            mean, var = self.moving_mean, self.moving_variance
 
-            conv_fold = tf.nn.conv2d(
-                self.inputs, W, strides=strides, padding=padding, use_cudnn_on_gpu=use_cudnn_on_gpu,
-                data_format=data_format
-            )
+        w_fold = self._w_fold(self.W, self.scale_para, var, self.epsilon)
 
-            self.outputs = tf.nn.bias_add(conv_fold, bias_fold, name='bn_bias_add')
+        W_ = quantize_weight_overflow(w_fold, self.bitW)
 
-            self.outputs = self._apply_activation(self.outputs)
+        conv_fold = tf.nn.conv2d(inputs, W_, strides=self.strides, padding=self.padding, data_format=self.data_format)
 
-        self._add_layers(self.outputs)
+        if self.beta_init:
+            bias_fold = self._bias_fold(self.offset_para, self.scale_para, mean, var, self.epsilon)
+            conv_fold = tf.nn.bias_add(conv_fold, bias_fold, name='bn_bias_add')
 
-        self._add_params([W, scale_para, offset_para, moving_mean, moving_variance])
+        if self.act:
+            conv_fold = self.act(conv_fold)
 
+        return conv_fold
 
-def _w_fold(w, gama, var, epsilon):
-    return tf.compat.v1.div(tf.multiply(gama, w), tf.sqrt(var + epsilon))
+    def mean_var_with_update(self, update_moving_mean, update_moving_variance, mean, variance):
+        with tf.control_dependencies([update_moving_mean, update_moving_variance]):
+            return tf.identity(mean), tf.identity(variance)
 
+    def _w_fold(self, w, gama, var, epsilon):
+        return tf.compat.v1.div(tf.multiply(gama, w), tf.sqrt(var + epsilon))
 
-def _bias_fold(beta, gama, mean, var, epsilon):
-    return tf.subtract(beta, tf.compat.v1.div(tf.multiply(gama, mean), tf.sqrt(var + epsilon)))
+    def _bias_fold(self, beta, gama, mean, var, epsilon):
+        return tf.subtract(beta, tf.compat.v1.div(tf.multiply(gama, mean), tf.sqrt(var + epsilon)))

tensorlayer/layers/dense/__init__.py

@@ -24,5 +24,5 @@
     'DropconnectDense',
     'TernaryDense',
     'QuanDense',
-    'QuanDenseLayerWithBN',
+    'QuanDenseWithBN',
 ]