joeyearsley / efficient_densenet_tensorflow Goto Github PK

class Back_Recompute(Layer):
    def __init__(self, nb_filter, bottleneck=False, dropout_rate=None, weight_decay=1e-4, **kwargs):
        self.nb_filter = nb_filter
        self.weight_decay = weight_decay
        self.bottleneck = bottleneck
        self.dropout_rate = dropout_rate
        super(Back_Recompute, self).__init__(**kwargs)

    def call(self, ip):
        global brcount
        concat_axis = 1 if K.image_data_format() == 'channels_first' else -1
        with tf.variable_scope('denseblock_{}'.format(brcount), use_resource=True):
            def _x(inner_ip):
                x = BatchNormalization(axis=concat_axis, epsilon=1.1e-5)(inner_ip)
                x = Activation('relu')(x)
                if self.bottleneck:
                    inter_channel = self.nb_filter * 4  # Obtained from https://github.com/liuzhuang13/DenseNet/blob/master/densenet.lua

                    x = Conv2D(inter_channel, (1, 1), kernel_initializer='he_normal', padding='same', use_bias=False,
                               kernel_regularizer=l2(self.weight_decay))(x)
                    x = BatchNormalization(axis=concat_axis, epsilon=1.1e-5)(x)
                    x = Activation('relu')(x)
                x = Conv2D(self.nb_filter, (3, 3), kernel_initializer='he_normal', padding='same', use_bias=False)(x)
                if self.dropout_rate:
                    x = Dropout(self.dropout_rate)(x)
                return x

            brcount = brcount + 1
            _x = tf.contrib.layers.recompute_grad(_x)

            return _x(ip)

    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1], input_shape[2], self.nb_filter)

ValueError: The variables used on recompute were different than the variables originally
used. The function wrapped with @recompute_grad likley creates its own variable
scope with a default name and has been called twice in the same enclosing scope.
To fix, ensure each call to the function happens in its own unique variable
scope.

def Conv2D_mem_eff( input_tensor,
                    filters,
                    kernel_size,
                    kernel_regularizer,
                    bias_regularizer,
                    padding,
                    name ):

    with tf.variable_scope( name,
                            use_resource = True ):

        def _x( inner_input_tensor ):

            x = Conv2D( filters = filters,
                        kernel_size = kernel_size,
                        padding = padding,
                        kernel_regularizer = kernel_regularizer,
                        bias_regularizer = bias_regularizer,
                        name = name )(inner_input_tensor)

            return x

        _x = tf.contrib.layers.recompute_grad( _x )

        return _x( input_tensor )

x = Lambda( Conv2D_mem_eff,
                     arguments = {'filters' : 24,
                                            'kernel_size' : (5,5),
                                            'kernel_regularizer' : l2,
                                            'bias_regularizer' : l2,
                                            'padding' : 'same',
                                            'name' : 'conv02'},
                     name= 'conv02' )(x)