├── .gitignore ├── README.md ├── main.py ├── module.py ├── net.py └── shuffle-block.png /.gitignore: -------------------------------------------------------------------------------- 1 | venv/* 2 | __pycache__/* 3 | .idea/* 4 | **/.DS_Store 5 | **/*.pyc 6 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Shufflenet-v2-tensorflow 2 | Tensorflow implementation of ECCV 2018 paper ShuffleNet V2. [[Paper]](https://arxiv.org/abs/1807.11164) 3 | 4 |

5 | 6 |

7 | 8 | # Support architectures 9 | * ShuffleNetV2 with 0.5, 1.0, 1.5, 2.0 channel multipliers 10 | 11 | # Usage 12 | ```python 13 | import tensorflow as tf 14 | from net import ShuffleNetV2 15 | 16 | if __name__ == '__main__': 17 | # Create input placeholder 18 | input = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3]) 19 | 20 | # Create model 21 | model = ShuffleNetV2(input, 1001, model_scale=1.0, is_training=True) 22 | 23 | # Get model inference result 24 | print(model.output) # shape = [None, 1001] 25 | 26 | ``` 27 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from net import ShuffleNetV2 3 | 4 | if __name__ == '__main__': 5 | # Create input placeholder 6 | input = tf.placeholder(dtype=tf.float32, shape=[None, 224, 224, 3]) 7 | 8 | # Create model 9 | model = ShuffleNetV2(input, 1001, model_scale=1.0, is_training=True) 10 | 11 | # Get model inference result 12 | print(model.output) # shape = [None, 1001] 13 | 14 | 15 | -------------------------------------------------------------------------------- /module.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | import tensorflow.contrib as tc 3 | 4 | slim = tc.slim 5 | 6 | def shuffle_unit(x, groups): 7 | with tf.variable_scope('shuffle_unit'): 8 | n, h, w, c = x.get_shape().as_list() 9 | x = tf.reshape(x, shape=tf.convert_to_tensor([tf.shape(x)[0], h, w, groups, c // groups])) 10 | x = tf.transpose(x, tf.convert_to_tensor([0, 1, 2, 4, 3])) 11 | x = tf.reshape(x, shape=tf.convert_to_tensor([tf.shape(x)[0], h, w, c])) 12 | return x 13 | 14 | 15 | def conv_bn_relu(x, out_channel, kernel_size, stride=1, dilation=1): 16 | with tf.variable_scope(None, 'conv_bn_relu'): 17 | x = slim.conv2d(x, out_channel, kernel_size, stride, rate=dilation, 18 | biases_initializer=None, activation_fn=None) 19 | x = slim.batch_norm(x, activation_fn=tf.nn.relu, fused=False) 20 | return x 21 | 22 | 23 | def conv_bn(x, out_channel, kernel_size, stride=1, dilation=1): 24 | with tf.variable_scope(None, 'conv_bn'): 25 | x = slim.conv2d(x, out_channel, kernel_size, stride, rate=dilation, 26 | biases_initializer=None, activation_fn=None) 27 | x = slim.batch_norm(x, activation_fn=None, fused=False) 28 | return x 29 | 30 | 31 | def depthwise_conv_bn(x, kernel_size, stride=1, dilation=1): 32 | with tf.variable_scope(None, 'depthwise_conv_bn'): 33 | x = slim.separable_conv2d(x, None, kernel_size, depth_multiplier=1, stride=stride, 34 | rate=dilation, activation_fn=None, biases_initializer=None) 35 | x = slim.batch_norm(x, activation_fn=None, fused=False) 36 | return x 37 | 38 | 39 | def resolve_shape(x): 40 | with tf.variable_scope(None, 'resolve_shape'): 41 | n, h, w, c = x.get_shape().as_list() 42 | if h is None or w is None: 43 | kernel_size = tf.convert_to_tensor([tf.shape(x)[1], tf.shape(x)[2]]) 44 | else: 45 | kernel_size = [h, w] 46 | return kernel_size 47 | 48 | 49 | def global_avg_pool2D(x): 50 | with tf.variable_scope(None, 'global_pool2D'): 51 | kernel_size = resolve_shape(x) 52 | x = slim.avg_pool2d(x, kernel_size, stride=1) 53 | x.set_shape([None, 1, 1, None]) 54 | return x 55 | 56 | 57 | def se_unit(x, bottleneck=2): 58 | with tf.variable_scope(None, 'SE_module'): 59 | n, h, w, c = x.get_shape().as_list() 60 | 61 | kernel_size = resolve_shape(x) 62 | x_pool = slim.avg_pool2d(x, kernel_size, stride=1) 63 | x_pool = tf.reshape(x_pool, shape=[-1, c]) 64 | fc = slim.fully_connected(x_pool, bottleneck, activation_fn=tf.nn.relu, 65 | biases_initializer=None) 66 | fc = slim.fully_connected(fc, c, activation_fn=tf.nn.sigmoid, 67 | biases_initializer=None) 68 | if n is None: 69 | channel_w = tf.reshape(fc, shape=tf.convert_to_tensor([tf.shape(x)[0], 1, 1, c])) 70 | else: 71 | channel_w = tf.reshape(fc, shape=[n, 1, 1, c]) 72 | 73 | x = tf.multiply(x, channel_w) 74 | return x 75 | 76 | 77 | def shufflenet_v2_block(x, out_channel, kernel_size, stride=1, dilation=1, shuffle_group=2): 78 | with tf.variable_scope(None, 'shuffle_v2_block'): 79 | if stride == 1: 80 | top, bottom = tf.split(x, num_or_size_splits=2, axis=3) 81 | 82 | half_channel = out_channel // 2 83 | 84 | top = conv_bn_relu(top, half_channel, 1) 85 | top = depthwise_conv_bn(top, kernel_size, stride, dilation) 86 | top = conv_bn_relu(top, half_channel, 1) 87 | 88 | out = tf.concat([top, bottom], axis=3) 89 | out = shuffle_unit(out, shuffle_group) 90 | 91 | else: 92 | half_channel = out_channel // 2 93 | b0 = conv_bn_relu(x, half_channel, 1) 94 | b0 = depthwise_conv_bn(b0, kernel_size, stride, dilation) 95 | b0 = conv_bn_relu(b0, half_channel, 1) 96 | 97 | b1 = depthwise_conv_bn(x, kernel_size, stride, dilation) 98 | b1 = conv_bn_relu(b1, half_channel, 1) 99 | 100 | out = tf.concat([b0, b1], axis=3) 101 | out = shuffle_unit(out, shuffle_group) 102 | return out 103 | 104 | 105 | -------------------------------------------------------------------------------- /net.py: -------------------------------------------------------------------------------- 1 | from module import * 2 | 3 | class ShuffleNetV2(): 4 | 5 | first_conv_channel = 24 6 | 7 | def __init__(self, input_holder, cls, model_scale=1.0, shuffle_group=2, is_training=True): 8 | self.input = input_holder 9 | self.output = None 10 | self.cls = cls 11 | self.shuffle_group = shuffle_group 12 | self.channel_sizes = self._select_channel_size(model_scale) 13 | 14 | with slim.arg_scope([slim.batch_norm], is_training=is_training): 15 | self._build_model() 16 | 17 | def _select_channel_size(self, model_scale): 18 | # [(out_channel, repeat_times), (out_channel, repeat_times), ...] 19 | if model_scale == 0.5: 20 | return [(48, 4), (96, 8), (192, 4), (1024, 1)] 21 | elif model_scale == 1.0: 22 | return [(116, 4), (232, 8), (464, 4), (1024, 1)] 23 | elif model_scale == 1.5: 24 | return [(176, 4), (352, 8), (704, 4), (1024, 1)] 25 | elif model_scale == 2.0: 26 | return [(244, 4), (488, 8), (976, 4), (2048, 1)] 27 | else: 28 | raise ValueError('Unsupported model size.') 29 | 30 | def _build_model(self): 31 | with tf.variable_scope('init_block'): 32 | out = conv_bn_relu(self.input, self.first_conv_channel, 3, 2) 33 | out = slim.max_pool2d(out, 3, 2, padding='SAME') 34 | 35 | for idx, block in enumerate(self.channel_sizes[:-1]): 36 | with tf.variable_scope('shuffle_block_{}'.format(idx)): 37 | out_channel, repeat = block 38 | 39 | # First block is downsampling 40 | out = shufflenet_v2_block(out, out_channel, 3, 2, shuffle_group=self.shuffle_group) 41 | 42 | # Rest blocks 43 | for i in range(repeat-1): 44 | out = shufflenet_v2_block(out, out_channel, 3, shuffle_group=self.shuffle_group) 45 | 46 | with tf.variable_scope('end_block'): 47 | out = conv_bn_relu(out, self.channel_sizes[-1][0], 1) 48 | 49 | with tf.variable_scope('prediction'): 50 | out = global_avg_pool2D(out) 51 | out = slim.conv2d(out, self.cls, 1, activation_fn=None, biases_initializer=None) 52 | out = tf.reshape(out, shape=[-1, self.cls]) 53 | out = tf.identity(out, name='cls_prediction') 54 | self.output = out 55 | 56 | -------------------------------------------------------------------------------- /shuffle-block.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/timctho/shufflenet-v2-tensorflow/ae091dfbf10e5bf0fb723e00ebbf5410b550f4f8/shuffle-block.png --------------------------------------------------------------------------------