shufflenet_v1_version1.py

# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
slim = tf.contrib.slim
import functools
import numpy as np

trunc_normal = lambda stddev: tf.truncated_normal_initializer(0.0, stddev)

def channel_shuffle(input, output, group, scope=None):
        assert 0 == output % group, "Output channels must be a multiple of groups"
        num_channels_in_group = output // group
        with tf.variable_scope(scope,"ChannelShuffle",[input]):
            net = tf.split(input, output, axis=3, name="split")
            chs = []
            for i in range(group):
                for j in range(num_channels_in_group):
                    chs.append(net[i + j * group])
            net = tf.concat(chs, axis=3, name="concat")
        return net
    
def channel_shuffle_v1(input,depth_bottleneck,group,scope=None):
    assert 0 == depth_bottleneck % group, "Output channels must be a multiple of groups"
    with tf.variable_scope(scope,"ChannelShuffle",[input]):
        n, h, w, c =input.shape.as_list()
        x_reshape = tf.reshape(input, [-1,h,w,group,depth_bottleneck//group])
        x_transposed =tf.transpose(x_reshape, [0,1,2,4,3])
        net = tf.reshape(x_transposed, [-1,h,w,c])
        return net
        
def group_pointwise_conv2d(inputs, depth, stride, group, relu=True, scope=None):
    
        assert 0 == depth % group, "Output channels must be a multiple of groups"
        num_channels_in_group = depth // group
        with tf.variable_scope(scope, 'GConv', [inputs]) as sc:
            net = tf.split(inputs, group, axis=3, name="split")
            for i in range(group):
                net[i] = slim.conv2d(net[i],
                                     num_channels_in_group,
                                     [1, 1],
                                     stride=stride,
                                     activation_fn=None,
                                     normalizer_fn=None)
            net = tf.concat(net, axis=3, name="concat")
            net = slim.batch_norm(net, activation_fn = tf.nn.relu if relu else None)    
        return net    
    
@slim.add_arg_scope
def shuffle_bottleneck(inputs, depth_bottleneck, group, stride, shuffle=True, groups_in=None, outputs_collections = None, scope= None):
        if 1 != stride:
            _b, _h, _w, _c = inputs.get_shape().as_list()
            depth_bottleneck = depth_bottleneck - _c

        assert 0 == depth_bottleneck % group, "Output channels must be a multiple of groups"

        with tf.variable_scope(scope, 'Unit', [inputs]) as sc:
            if 1 != stride:
                net_skip = slim.avg_pool2d(inputs, [3, 3], stride, padding="SAME", scope='3x3AVGPool2D')
            else:
                net_skip = inputs

            net = group_pointwise_conv2d(inputs, depth_bottleneck, 1, group = (1 if groups_in is None else group), relu=True, scope="1x1GConvIn")

            if shuffle:
                net = channel_shuffle_v1(net, depth_bottleneck, group)


            # separable_conv2d produces only a depthwise convolution layer
            net = slim.separable_conv2d(net, None, [3, 3],
                                        depth_multiplier=1,
                                        stride=stride,
                                        rate=1,
                                        normalizer_fn=slim.batch_norm,
                                        activation_fn=None,
                                        scope="3x3DWConv")

            net = group_pointwise_conv2d(net, depth_bottleneck, 1, group,relu=False, scope="1x1GConvOut")
            

            if 1 != stride:
                net = tf.concat([net, net_skip], axis=3)
            else:
                net = net + net_skip
            out = tf.nn.relu(net)
        return slim.utils.collect_named_outputs(outputs_collections, sc.name, out)
    
   
def shuffle_stage(inputs, depth, groups, repeat, shuffle=True,  groups_in=None,scope=None):
     with tf.variable_scope(scope,"Stage",[inputs]) as sc:
        net = shuffle_bottleneck(inputs, depth, group = groups, stride = 2, shuffle=shuffle ,groups_in=groups_in, scope='Unit{}'.format(0))
        
        for i in range(repeat):
            net = shuffle_bottleneck(net, depth, group = groups, stride = 1, shuffle=shuffle , scope='Unit{}'.format(i + 1))   
        return net   

def shufflenet(inputs, 
               num_classes=1000,
               dropout_keep_prob=0.90, 
               is_training=True,
               shuffle=True, 
               base_ch=144, 
               groups=1, 
               prediction_fn=tf.contrib.layers.softmax,
               spatial_squeeze=True,
               reuse=None,
               scope='ShuffleNet',
               global_pool=True):
    input_shape = inputs.get_shape().as_list()
    if len(input_shape) != 4:
        raise ValueError('Invalid input tensor rank, expected 4, was: %d' %
                     len(input_shape))
          
    with tf.variable_scope(scope, 'ShuffleNet', [inputs], reuse=reuse) as sc:
        end_points_collection = sc.original_name_scope + '_end_points' 
        with slim.arg_scope([slim.conv2d,slim.separable_conv2d,slim.avg_pool2d,slim.max_pool2d,shuffle_bottleneck],
                            outputs_collections = [end_points_collection]):
            with slim.arg_scope([slim.batch_norm, slim.dropout], is_training=is_training):
            
                with tf.variable_scope('Stage1'):
                    net = slim.conv2d(inputs, 24, [3, 3], stride = 2, scope= "conv1")
                    net = slim.max_pool2d(net, [3, 3], stride = 2, padding='SAME', scope= "pool1")
        
                net = shuffle_stage(net, depth = base_ch * 1, groups = groups, repeat = 3, groups_in = 1, scope='Stage2')
                net = shuffle_stage(net, depth = base_ch * 2, groups = groups, repeat = 7, scope='Stage3')
                net = shuffle_stage(net, depth = base_ch * 4, groups = groups, repeat = 3, scope='Stage4')
    
                
                # Convert end_points_collection into a dictionary of end_points.
                end_points = slim.utils.convert_collection_to_dict(end_points_collection)
        
                with tf.variable_scope('Logits'):      
                    if global_pool:
                        # Global average pooling.
                        net = tf.reduce_mean(net, [1, 2], name='global_pool', keep_dims=True)
                        end_points['global_pool'] = net
                    if not num_classes:
                        return net, end_points
                    # 1 x 1 x channel_dimensions
                    net = slim.dropout(net, keep_prob=dropout_keep_prob, scope='Dropout_1b')
                    logits = slim.conv2d(net, num_classes, [1, 1], activation_fn=None,
                                        normalizer_fn=None, scope='logits')
                    if spatial_squeeze:
                        logits = tf.squeeze(logits, [1, 2], name='SpatialSqueeze')
                
                end_points['Logits'] = logits
                if prediction_fn:
                    end_points['predictions'] = prediction_fn(logits, scope='predictions')

            return logits, end_points

shufflenet.default_image_size = 224    


def wrapped_partial(func, *args, **kwargs):
  partial_func = functools.partial(func, *args, **kwargs)
  functools.update_wrapper(partial_func, func)
  return partial_func

shufflenet_g1 = wrapped_partial(shufflenet, shuffle=False, base_ch=144, groups=1)
shufflenet_g2 = wrapped_partial(shufflenet, base_ch=200, groups=2)
shufflenet_g3 = wrapped_partial(shufflenet, base_ch=240, groups=3)
shufflenet_g4 = wrapped_partial(shufflenet, base_ch=272, groups=4)
shufflenet_g8 = wrapped_partial(shufflenet, base_ch=384, groups=8)


def shufflenet_arg_scope(is_training = True,
                    weight_decay = 0.0001,
                    batch_norm_decay = 0.997,
                    batch_norm_epsilon = 1e-5,
                    batch_norm_scale = True):

    batch_norm_params = {
        'is_training': is_training,
        'decay': batch_norm_decay,
        'epsilon': batch_norm_epsilon,
        'scale': batch_norm_scale,
        'updates_collections': tf.GraphKeys.UPDATE_OPS,
    }

    with slim.arg_scope(
            [slim.conv2d],
            weights_regularizer = slim.l2_regularizer(weight_decay),
            weights_initializer = slim.variance_scaling_initializer(),
            activation_fn = tf.nn.relu,
            normalizer_fn = slim.batch_norm,
            normalizer_params = batch_norm_params):
        with slim.arg_scope([slim.batch_norm], **batch_norm_params):
            with slim.arg_scope([slim.max_pool2d], padding = 'SAME') as arg_sc:
                return arg_sc



if __name__ == "__main__":
    inputs = tf.random_normal([1, 224, 224, 3])
    
    logits, end_points= shufflenet_g2(inputs,num_classes=1000)
    print(end_points['predictions'])   
    writer = tf.summary.FileWriter("./logs", graph=tf.get_default_graph())
    
    print("Layers")
    for k, v in end_points.items():
        print('name = {}, shape = {}'.format(v.name, v.get_shape()))
  
    print("Parameters")
    for v in slim.get_model_variables():
        print('name = {}, shape = {}'.format(v.name, v.get_shape()))  
    
    
    
    init = tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init)
        pred = sess.run(end_points['predictions'])
#         print(pred)
        print(np.argmax(pred,1))
        print(pred[:,np.argmax(pred,1)])