def DynamicConvFilter(inputs,
                      filters,
                      out_channel,
                      kernel_shape,
                      stride=1,
                      padding='SAME'):
    """ see "Dynamic Filter Networks" (NIPS 2016)
        by Bert De Brabandere*, Xu Jia*, Tinne Tuytelaars and Luc Van Gool

    Remarks:
        This is the convolution version of a dynamic filter.

    Args:
        inputs : unfiltered input [b, h, w, 1] only grayscale images.
        filters : learned filters of [b, k, k, 1] (dynamically generated by the network).
        out_channel (int): number of output channel.
        kernel_shape: (h, w) tuple or a int.
        stride: (h, w) tuple or a int.
        padding (str): 'valid' or 'same'. Case insensitive.

    Returns
        tf.Tensor named ``output``.
    """

    # tf.unstack only works with known batch_size :-(
    batch_size, h, w, in_channel = inputs.get_shape().as_list()
    stride = shape4d(stride)

    inputs = tf.unstack(inputs)
    filters = tf.reshape(filters, [batch_size] + shape2d(kernel_shape) +
                         [in_channel, out_channel])
    filters = tf.unstack(filters)

    # this is ok as TF uses the cuda stream context
    rsl = [
        tf.nn.conv2d(
            tf.reshape(d, [1, h, w, in_channel]),
            tf.reshape(k,
                       [kernel_shape, kernel_shape, in_channel, out_channel]),
            stride,
            padding="SAME") for d, k in zip(inputs, filters)
    ]
    rsl = tf.concat(rsl, axis=0, name='output')
    return rsl
Esempio n. 2
0
def DynamicConvFilter(inputs, filters, out_channel,
                      kernel_shape,
                      stride=1,
                      padding='SAME'):
    """ see "Dynamic Filter Networks" (NIPS 2016)
        by Bert De Brabandere*, Xu Jia*, Tinne Tuytelaars and Luc Van Gool

    Remarks:
        This is the convolution version of a dynamic filter.

    Args:
        inputs : unfiltered input [b, h, w, 1] only grayscale images.
        filters : learned filters of [b, k, k, 1] (dynamically generated by the network).
        out_channel (int): number of output channel.
        kernel_shape: (h, w) tuple or a int.
        stride: (h, w) tuple or a int.
        padding (str): 'valid' or 'same'. Case insensitive.

    Returns
        tf.Tensor named ``output``.
    """

    # tf.unstack only works with known batch_size :-(
    batch_size, h, w, in_channel = inputs.get_shape().as_list()
    stride = shape4d(stride)

    inputs = tf.unstack(inputs)
    filters = tf.reshape(filters, [batch_size] + shape2d(kernel_shape) + [in_channel, out_channel])
    filters = tf.unstack(filters)

    # this is ok as TF uses the cuda stream context
    rsl = [tf.nn.conv2d(tf.reshape(d, [1, h, w, in_channel]),
                        tf.reshape(k, [kernel_shape, kernel_shape, in_channel, out_channel]),
                        stride, padding="SAME") for d, k in zip(inputs, filters)]
    rsl = tf.concat(rsl, axis=0, name='output')
    return rsl
def Conv2DWithTrackedMults(x,
                           out_channel,
                           kernel_shape,
                           network_complexity=None,
                           padding='SAME',
                           stride=1,
                           W_init=None,
                           b_init=None,
                           nl=tf.identity,
                           split=1,
                           use_bias=True,
                           data_format='NHWC'):
    """
    2D convolution on 4D inputs.
    Args:
        x (tf.Tensor): a 4D tensor.
            Must have known number of channels, but can have other unknown dimensions.
        out_channel (int): number of output channel.
        kernel_shape: (h, w) tuple or a int.
        stride: (h, w) tuple or a int.
        padding (str): 'valid' or 'same'. Case insensitive.
        split (int): Split channels as used in Alexnet. Defaults to 1 (no split).
        W_init: initializer for W. Defaults to `variance_scaling_initializer`.
        b_init: initializer for b. Defaults to zero.
        nl: a nonlinearity function.
        use_bias (bool): whether to use bias.
    Returns:
        tf.Tensor named ``output`` with attribute `variables`.
    Variable Names:
    * ``W``: weights
    * ``b``: bias
    """
    in_shape = x.get_shape().as_list()
    channel_axis = 3 if data_format == 'NHWC' else 1
    in_channel = in_shape[channel_axis]
    assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!"
    assert in_channel % split == 0
    assert out_channel % split == 0

    kernel_shape = shape2d(kernel_shape)
    padding = padding.upper()
    filter_shape = kernel_shape + [in_channel / split, out_channel]
    stride = shape4d(stride, data_format=data_format)

    if W_init is None:
        W_init = tf.contrib.layers.variance_scaling_initializer()
    if b_init is None:
        b_init = tf.constant_initializer()

    W = tf.get_variable('W', filter_shape, initializer=W_init)
    network_complexity['weights'] += filter_shape[0] * filter_shape[
        1] * filter_shape[2] * filter_shape[3]

    if use_bias:
        b = tf.get_variable('b', [out_channel], initializer=b_init)
        network_complexity['weights'] += out_channel

    assert split == 1
    xsh = x.get_shape().as_list()
    network_complexity['mults'] += xsh[1] * xsh[2] * filter_shape[
        0] * filter_shape[1] * filter_shape[2] * filter_shape[3] / (stride[2] *
                                                                    stride[2])
    conv = tf.nn.conv2d(x, W, stride, padding, data_format=data_format)

    ret = nl(
        tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv,
        name='output')
    ret.variables = VariableHolder(W=W)
    if use_bias:
        ret.variables.b = b

    return ret
def Conv3D(
        inputs,
        filters,
        kernel_size,
        strides=(1, 1, 1),
        padding='same',
        data_format='channels_last',
        dilation_rate=(1, 1, 1),
        activation=None,
        use_bias=True,
        kernel_initializer=tf.contrib.layers.variance_scaling_initializer(2.0),
        bias_initializer=tf.zeros_initializer(),
        kernel_regularizer=None,
        bias_regularizer=None,
        activity_regularizer=None,
        split=1):
    """
    A wrapper around `tf.layers.Conv2D`.
    Some differences to maintain backward-compatibility:
    1. Default kernel initializer is variance_scaling_initializer(2.0).
    2. Default padding is 'same'.
    3. Support 'split' argument to do group conv.
    Variable Names:
    * ``W``: weights
    * ``b``: bias
    """
    if split == 1:
        with rename_get_variable({'kernel': 'W', 'bias': 'b'}):
            layer = tf.layers.Conv3D(filters,
                                     kernel_size,
                                     strides=strides,
                                     padding=padding,
                                     data_format=data_format,
                                     dilation_rate=dilation_rate,
                                     activation=activation,
                                     use_bias=use_bias,
                                     kernel_initializer=kernel_initializer,
                                     bias_initializer=bias_initializer,
                                     kernel_regularizer=kernel_regularizer,
                                     bias_regularizer=bias_regularizer,
                                     activity_regularizer=activity_regularizer)
            ret = layer.apply(inputs, scope=tf.get_variable_scope())
            ret = tf.identity(ret, name='output')

        ret.variables = VariableHolder(W=layer.kernel)
        if use_bias:
            ret.variables.b = layer.bias

    else:
        # group conv implementation
        data_format = get_data_format(data_format, tfmode=False)
        in_shape = inputs.get_shape().as_list()
        channel_axis = 3 if data_format == 'NHWC' else 1
        in_channel = in_shape[channel_axis]
        assert in_channel is not None, "[Conv3D] Input cannot have unknown channel!"
        assert in_channel % split == 0

        assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
            "Not supported by group conv now!"

        out_channel = filters
        assert out_channel % split == 0
        assert dilation_rate == (1, 1) or get_tf_version_number(
        ) >= 1.5, 'TF>=1.5 required for group dilated conv'

        kernel_shape = shape2d(kernel_size)
        filter_shape = kernel_shape + [in_channel / split, out_channel]
        stride = shape4d(strides, data_format=data_format)

        kwargs = dict(data_format=data_format)
        if get_tf_version_number() >= 1.5:
            kwargs['dilations'] = shape4d(dilation_rate,
                                          data_format=data_format)

        W = tf.get_variable('W', filter_shape, initializer=kernel_initializer)

        if use_bias:
            b = tf.get_variable('b', [out_channel],
                                initializer=bias_initializer)

        inputs = tf.split(inputs, split, channel_axis)
        kernels = tf.split(W, split, 3)
        outputs = [
            tf.nn.conv2d(i, k, stride, padding.upper(), **kwargs)
            for i, k in zip(inputs, kernels)
        ]
        conv = tf.concat(outputs, channel_axis)
        if activation is None:
            activation = tf.identity
        ret = activation(tf.nn.bias_add(conv, b, data_format=data_format)
                         if use_bias else conv,
                         name='output')

        ret.variables = VariableHolder(W=W)
        if use_bias:
            ret.variables.b = b
    return ret
Esempio n. 5
0
def Conv(inputs,
         filters,
         kernel_size,
         strides=(1, 1),
         padding='same',
         data_format='channels_last',
         dilation_rate=(1, 1),
         activation=None,
         use_bias=True,
         kernel_initializer=None,
         bias_initializer=tf.zeros_initializer(),
         kernel_regularizer=None,
         bias_regularizer=None,
         activity_regularizer=None,
         split=1,
         norm=False):
    """
    Similar to `tf.layers.Conv2D`, but with some differences:
    1. Default kernel initializer is variance_scaling_initializer(2.0).
    2. Default padding is 'same'.
    3. Support 'split' argument to do group convolution.
    Variable Names:
    * ``W``: weights
    * ``b``: bias
    """
    if kernel_initializer is None:
        if get_tf_version_tuple() <= (1, 12):
            kernel_initializer = tf.contrib.layers.variance_scaling_initializer(
                2.0)  # deprecated
        else:
            kernel_initializer = tf.keras.initializers.VarianceScaling(
                2.0, distribution='untruncated_normal')
    dilation_rate = shape2d(dilation_rate)

    if True:
        # group conv implementation
        data_format = get_data_format(data_format, keras_mode=False)
        in_shape = inputs.get_shape().as_list()
        channel_axis = 3 if data_format == 'NHWC' else 1
        in_channel = in_shape[channel_axis]
        assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!"
        assert in_channel % split == 0

        assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
            "Not supported by group conv or dilated conv!"

        out_channel = filters
        assert out_channel % split == 0
        assert dilation_rate == [1, 1] or get_tf_version_tuple() >= (
            1, 5), 'TF>=1.5 required for dilated conv.'

        kernel_shape = shape2d(kernel_size)
        filter_shape = kernel_shape + [in_channel // split, out_channel]
        stride = shape4d(strides, data_format=data_format)

        kwargs = {"data_format": data_format}
        if get_tf_version_tuple() >= (1, 5):
            kwargs['dilations'] = shape4d(dilation_rate,
                                          data_format=data_format)

        # matching input dtype (ex. tf.float16) since the default dtype of variable if tf.float32
        inputs_dtype = inputs.dtype
        W = tf.get_variable('parseweigth',
                            filter_shape,
                            dtype=inputs_dtype,
                            initializer=kernel_initializer)
        if norm:
            use_bias = False
            W = tf.reshape(W, kernel_shape + [4, in_channel // 4, out_channel])
            W = tf.nn.softmax(W, 2)
            W = tf.reshape(W, filter_shape)
        #dynamics = tf.reduce_mean(inputs, 0)
        #dynamics = tf.transpose(dynamics, [1,2,0])
        #dynamics = tf.image.resize_images(dynamics, kernel_shape)
        #dynamics = tf.expand_dims(dynamics, -1)
        #W = W  +  0.001 * dynamics #tf.random_normal(shape = tf.shape(W), mean = 0.0, stddev = 0.012, dtype = tf.float32)

        #W = W *tf.random_uniform(shape=W.get_shape().as_list(), minval=0., maxval=2.)

        if use_bias:
            b = tf.get_variable('parsebias', [out_channel],
                                dtype=inputs_dtype,
                                initializer=bias_initializer)

        if split == 1:
            conv = tf.nn.conv2d(inputs, W, stride, padding.upper(), **kwargs)
        else:
            try:
                conv = tf.nn.conv2d(inputs, W, stride, padding.upper(),
                                    **kwargs)
            except ValueError:
                log_once(
                    "CUDNN group convolution support is only available with "
                    "https://github.com/tensorflow/tensorflow/pull/25818 . "
                    "Will fall back to a loop-based slow implementation instead!",
                    'warn')

        ret = tf.nn.bias_add(conv, b,
                             data_format=data_format) if use_bias else conv
        if activation is not None:
            ret = activation(ret)
        ret = tf.identity(ret, name='output')

        ret.variables = VariableHolder(W=W)
        if use_bias:
            ret.variables.b = b
    return ret
Esempio n. 6
0
def mpusim_depthwise_convolution2d(inputs,
                                   kernel_size,
                                   strides=(1, 1),
                                   padding='valid',
                                   depth_multiplier=1,
                                   data_format='channels_last',
                                   activation=None,
                                   use_bias=True,
                                   depthwise_initializer='glorot_uniform',
                                   bias_initializer='zeros',
                                   depthwise_regularizer=None,
                                   bias_regularizer=None,
                                   depthwise_constraint=None,
                                   bias_constraint=None,
                                   activations_datatype_size_byte=1,
                                   weights_datatype_size_byte=1,
                                   results_datatype_size_byte=4,
                                   systolic_array_height=256,
                                   systolic_array_width=256,
                                   activation_fifo_depth=8,
                                   accumulator_array_height=4096,
                                   log_file_output_dir='.',
                                   model_name='unnamed'):

    #depthwise_initializer = initializers.get(depthwise_initializer)
    #depthwise_regularizer = regularizers.get(depthwise_regularizer)
    #depthwise_constraint = constraints.get(depthwise_constraint)
    #bias_initializer = initializers.get(bias_initializer)

    data_format = get_data_format(data_format, keras_mode=False)
    input_shape = inputs.get_shape().as_list()

    strides = shape4d(strides, data_format=data_format)

    if len(input_shape) < 4:
        raise ValueError(
            'Inputs to `mpusim_depthwise_conv2d` should have rank 4. '
            'Received input shape:', str(input_shape))

    if data_format == 'NCHW':
        raise ValueError('mpusim_depthwise_convolution2d '
                         'only supports NHWC data format')
    else:
        channel_axis = 3

    if input_shape[channel_axis] is None:
        raise ValueError('The channel dimension of the inputs to '
                         '`mpusim_depthwise_convolution2d` '
                         'should be defined. Found `None`.')

    input_dim = int(input_shape[channel_axis])

    depthwise_kernel_shape = (kernel_size[0], kernel_size[1], input_dim,
                              depth_multiplier)

    depthwise_kernel = tf.get_variable('W',
                                       shape=depthwise_kernel_shape,
                                       initializer=depthwise_initializer,
                                       regularizer=depthwise_regularizer,
                                       constraint=depthwise_constraint)

    if use_bias:
        biases = tf.get_variable('b',
                                 shape=(input_dim * depth_multiplier, ),
                                 initializer=bias_initializer,
                                 regularizer=bias_regularizer,
                                 constraint=bias_constraint)

    result = mpusim_depthwise_conv2d(
        inputs,
        depthwise_kernel,
        strides=strides,
        padding=padding,
        data_format=data_format,
        activations_datatype_size_byte=activations_datatype_size_byte,
        weights_datatype_size_byte=weights_datatype_size_byte,
        results_datatype_size_byte=results_datatype_size_byte,
        systolic_array_height=systolic_array_height,
        systolic_array_width=systolic_array_width,
        activation_fifo_depth=activation_fifo_depth,
        accumulator_array_height=accumulator_array_height,
        log_file_output_dir=log_file_output_dir,
        model_name=model_name)

    if use_bias:
        result = tf.nn.bias_add(result, bias, data_format=data_format)

    if activation is not None:
        result = activation(result)

    result = tf.identity(result, name='output')

    result.variables = VariableHolder(W=depthwise_kernel)

    if use_bias:
        result.variables.b = biases

    return result
Esempio n. 7
0
def MaskedConv2D(
        inputs,
        filters,
        kernel_size,
        strides=(1, 1),
        padding='same',
        data_format='channels_last',
        dilation_rate=(1, 1),
        activation=None,
        use_bias=True,
        kernel_initializer=None,
        bias_initializer=tf.zeros_initializer(),
        kernel_regularizer=None,
        bias_regularizer=None,
        activity_regularizer=None,
        split=1,
        masking=False):
    """
    A wrapper around `tf.layers.Conv2D`.
    Some differences to maintain backward-compatibility:

    1. Default kernel initializer is variance_scaling_initializer(2.0).
    2. Default padding is 'same'.
    3. Support 'split' argument to do group conv.

    Variable Names:

    * ``W``: weights
    * ``b``: bias
    """
    if kernel_initializer is None:
        if get_tf_version_tuple() <= (1, 12):
            kernel_initializer = tf.contrib.layers.variance_scaling_initializer(2.0)
        else:
            kernel_initializer = tf.keras.initializers.VarianceScaling(2.0, distribution='untruncated_normal')
    dilation_rate = shape2d(dilation_rate)

    if (masking == False) and (split == 1) and (dilation_rate == [1, 1]):
        # tf.layers.Conv2D has bugs with dilations (https://github.com/tensorflow/tensorflow/issues/26797)
        with rename_get_variable({'kernel': 'W', 'bias': 'b'}):
            layer = tf.layers.Conv2D(
                filters,
                kernel_size,
                strides=strides,
                padding=padding,
                data_format=data_format,
                dilation_rate=dilation_rate,
                activation=activation,
                use_bias=use_bias,
                kernel_initializer=kernel_initializer,
                bias_initializer=bias_initializer,
                kernel_regularizer=kernel_regularizer,
                bias_regularizer=bias_regularizer,
                activity_regularizer=activity_regularizer,
                _reuse=tf.get_variable_scope().reuse)
            ret = layer.apply(inputs, scope=tf.get_variable_scope())
            ret = tf.identity(ret, name='output')

        ret.variables = VariableHolder(W=layer.kernel)
        if use_bias:
            ret.variables.b = layer.bias

    else:
        if masking == True:
            assert split == 1, "Pruining group conv is not supported yet"

        # group conv implementation
        data_format = get_data_format(data_format, keras_mode=False)
        in_shape = inputs.get_shape().as_list()
        channel_axis = 3 if data_format == 'NHWC' else 1
        in_channel = in_shape[channel_axis]
        assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!"
        assert in_channel % split == 0

        assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
            "Not supported by group conv or dilated conv!"

        out_channel = filters
        assert out_channel % split == 0
        assert dilation_rate == [1, 1] or get_tf_version_tuple() >= (1, 5), 'TF>=1.5 required for dilated conv.'

        kernel_shape = shape2d(kernel_size)
        filter_shape = kernel_shape + [in_channel / split, out_channel]
        stride = shape4d(strides, data_format=data_format)

        kwargs = dict(data_format=data_format)
        if get_tf_version_tuple() >= (1, 5):
            kwargs['dilations'] = shape4d(dilation_rate, data_format=data_format)

        W = tf.get_variable(
            'W', filter_shape, initializer=kernel_initializer)

        if use_bias:
            b = tf.get_variable('b', [out_channel], initializer=bias_initializer)

        if split == 1:
            if masking:
                W = pruning.apply_mask(W)
            conv = tf.nn.conv2d(inputs, W, stride, padding.upper(), **kwargs)
        else:
            conv = None
            if get_tf_version_tuple() >= (1, 13):
                try:
                    conv = tf.nn.conv2d(inputs, W, stride, padding.upper(), **kwargs)
                except ValueError:
                    log_once("CUDNN group convolution support is only available with "
                             "https://github.com/tensorflow/tensorflow/pull/25818 . "
                             "Will fall back to a loop-based slow implementation instead!", 'warn')
            if conv is None:
                inputs = tf.split(inputs, split, channel_axis)
                kernels = tf.split(W, split, 3)
                outputs = [tf.nn.conv2d(i, k, stride, padding.upper(), **kwargs)
                           for i, k in zip(inputs, kernels)]
                conv = tf.concat(outputs, channel_axis)

        ret = tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv
        if activation is not None:
            ret = activation(ret)
        ret = tf.identity(ret, name='output')

        ret.variables = VariableHolder(W=W)
        if use_bias:
            ret.variables.b = b
    return ret
Esempio n. 8
0
def mpusim_conv2d(
        inputs,
        filters,
        kernel_size,
        strides=(1, 1),
        padding='same',
        data_format='channels_last',
        dilation_rate=(1, 1),
        activation=None,
        use_bias=True,
        kernel_initializer=None,
        bias_initializer=tf.zeros_initializer(),
        kernel_regularizer=None,
        bias_regularizer=None,
        activity_regularizer=None,
        split=1,
        activations_datatype_size_byte=1,
        weights_datatype_size_byte=1,
        results_datatype_size_byte=4,
        systolic_array_height=256,
        systolic_array_width=256,
        activation_fifo_depth=8,
        accumulator_array_height=4096,
        log_file_output_dir='.',
        model_name='unnamed'):
    """
    Similar to `tf.layers.Conv2D`, but with some differences:

    1. Default kernel initializer is variance_scaling_initializer(2.0).
    2. Default padding is 'same'.
    3. Support 'split' argument to do group convolution.

    Variable Names:

    * ``W``: weights
    * ``b``: bias
    """
    if kernel_initializer is None:
        if get_tf_version_tuple() <= (1, 12):
            kernel_initializer = tf.contrib.layers.variance_scaling_initializer(2.0)
        else:
            kernel_initializer = tf.keras.initializers.VarianceScaling(2.0, distribution='untruncated_normal')
    dilation_rate = shape2d(dilation_rate)

    # group conv implementation
    data_format = get_data_format(data_format, keras_mode=False)
    in_shape = inputs.get_shape().as_list()
    channel_axis = 3 if data_format == 'NHWC' else 1
    in_channel = in_shape[channel_axis]
    assert in_channel is not None, "[mpusim_conv2d] Input cannot have unknown channel!"
    assert in_channel % split == 0

    assert kernel_regularizer is None and bias_regularizer is None and activity_regularizer is None, \
        "Not supported by group conv or dilated conv!"

    out_channel = filters
    assert out_channel % split == 0
    assert dilation_rate == [1, 1] or get_tf_version_tuple() >= (1, 5), 'TF>=1.5 required for dilated conv.'

    kernel_shape = shape2d(kernel_size)
    filter_shape = kernel_shape + [in_channel / split, out_channel]
    stride = shape4d(strides, data_format=data_format)

    kwargs = dict(data_format=data_format)
    if get_tf_version_tuple() >= (1, 5):
        kwargs['dilations'] = shape4d(dilation_rate, data_format=data_format)

    W = tf.get_variable(
            'W', filter_shape, initializer=kernel_initializer)

    if use_bias:
        b = tf.get_variable('b', [out_channel], initializer=bias_initializer)

    if split == 1:
        conv = mpu_sim_conv2d_lib.mpu_sim_conv2d(inputs,
                                                    W,
                                                    activations_datatype_size_byte,
                                                    weights_datatype_size_byte,
                                                    results_datatype_size_byte,
                                                    systolic_array_height,
                                                    systolic_array_width,
                                                    activation_fifo_depth,
                                                    accumulator_array_height,
                                                    log_file_output_dir,
                                                    model_name,
                                                    stride,
                                                    padding.upper(),
                                                    **kwargs)
    else:
        
        inputs = tf.split(inputs, split, channel_axis)
        kernels = tf.split(W, split, 3)
        outputs = [mpu_sim_conv2d_lib.mpu_sim_conv2d(input_block,
                                                        kernel_block,
                                                        activations_datatype_size_byte,
                                                        weights_datatype_size_byte,
                                                        results_datatype_size_byte,
                                                        systolic_array_height,
                                                        systolic_array_width,
                                                        activation_fifo_depth,
                                                        accumulator_array_height,
                                                        log_file_output_dir,
                                                        model_name,
                                                        stride,
                                                        padding.upper(),
                                                        **kwargs)
                    for input_block, kernel_block in zip(inputs, kernels)]
        conv = tf.concat(outputs, channel_axis)

    ret = tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv
    if activation is not None:
        ret = activation(ret)
    ret = tf.identity(ret, name='output')

    ret.variables = VariableHolder(W=W)
    if use_bias:
        ret.variables.b=b
    return ret
Esempio n. 9
0
def MyConv2D(x,
             out_channel,
             kernel_shape,
             padding='SAME',
             stride=1,
             W_init=None,
             b_init=None,
             activation=tf.identity,
             split=1,
             use_bias=True,
             data_format='NHWC',
             wscale=True):

    in_shape = x.get_shape().as_list()
    channel_axis = 3 if data_format == 'NHWC' else 1
    in_channel = in_shape[channel_axis]
    assert in_channel is not None, "[Conv2D] Input cannot have unknown channel!"
    assert in_channel % split == 0
    assert out_channel % split == 0

    kernel_shape = shape2d(kernel_shape)
    padding = padding.upper()
    filter_shape = kernel_shape + [in_channel / split, out_channel]
    stride = shape4d(stride, data_format=data_format)

    # see paper details -- START
    fan_in = kernel_shape[0] * kernel_shape[1] * in_channel
    c = tf.constant(np.sqrt(2. / fan_in), dtype=tf.float32)

    # "use a trivial N(0,1)"
    W_init = tf.random_normal_initializer(stddev=c)
    W = tf.get_variable('W', filter_shape, initializer=W_init)

    # "scale weights at runtime"
    scale = tf.sqrt(tf.reduce_mean(W**2))

    W = W / scale

    if b_init is None:
        b_init = tf.constant_initializer()

    if use_bias:
        b = tf.get_variable('b', [out_channel], initializer=b_init)

    if split == 1:
        conv = tf.nn.conv2d(x, W, stride, padding, data_format=data_format)
    else:
        inputs = tf.split(x, split, channel_axis)
        kernels = tf.split(W, split, 3)
        outputs = [
            tf.nn.conv2d(i, k, stride, padding, data_format=data_format)
            for i, k in zip(inputs, kernels)
        ]
        conv = tf.concat(outputs, channel_axis)

    conv = conv * scale

    ret = activation(
        tf.nn.bias_add(conv, b, data_format=data_format) if use_bias else conv,
        name='output')
    ret.variables = VariableHolder(W=W)
    if use_bias:
        ret.variables.b = b
    return ret