add train method.

README.md

@@ -1,11 +1,11 @@
 # MobileNet v2 
-A Keras 2 implementation of MobileNet V2.  
+A Python 3 and Keras 2 implementation of MobileNet V2 and provide train method.  
 
-According to the paper:[Inverted Residuals and Linear Bottlenecks Mobile Networks for Classification, Detection and Segmentation](https://arxiv.org/abs/1801.04381)
+According to the paper: [Inverted Residuals and Linear Bottlenecks Mobile Networks for Classification, Detection and Segmentation](https://arxiv.org/abs/1801.04381).
 
-Currently only the network structure is defined, and the training function will be updated later.
 
 ## Requirement
+- OpenCV 3.4
 - Python 3.5    
 - Tensorflow-gpu 1.2.0  
 - Keras 2.1.3
@@ -22,14 +22,56 @@ Each line describes a sequence of 1 or more identical (modulo stride) layers, re
 
 ![residual block architectures](/images/stru.jpg)
 
-**Architectures of this implementation with (224, 224, 3) inputs and 1000 output:**
 
-![architectures](/images/MobileNetv2.png)
+## Train the model
 
-##Reference
-- [Inverted Residuals and Linear Bottlenecks Mobile Networks for Classification, Detection and Segmentation](https://arxiv.org/abs/1801.04381)
+The recommended size of the image in the paper is 224 * 224. The ```data\convert.py``` file provide a demo of resize cifar-100 dataset to this size.
 
-##Copyright
+**The dataset folder structure is as follows:**
+
+	data
+	--train
+	  --class 0
+	  --class 1
+	  ....
+	  --class n
+	--validation
+	  --class 0
+	  --class 1
+	  ....
+	  --class n
+
+**Run command below to train the model:**
+
+```
+python train.py --classes num_classes --batch batch_size --epochs epochs --size image_size
+```
+
+The ```.h5``` weight file was saved at model folder. If you want to do fine tune the trained model, you can run the following command. However, it should be noted that the size of the input image should be consistent with the original model.
+
+```
+python train.py --classes num_classes --batch batch_size --epochs epochs --size image_size --weights weights_path
+```
+
+**Parameter explanation**
+
+- --classes, The number of classes of dataset.  
+- --size,    The image size of train sample.  
+- --batch,   The number of train samples per batch.  
+- --epochs,  The number of train iterations.  
+- --weights, Fine tune with other weights.  
+
+## Reference
+
+	@article{MobileNetv2,  
+	  title={Inverted Residuals and Linear Bottlenecks Mobile Networks for Classification, Detection and Segmentatio},  
+	  author={Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen},
+	  journal={arXiv preprint arXiv:1801.04381},
+	  year={2018}
+	}
+
+
+## Copyright
 See [LICENSE](LICENSE) for details.
 
 

data/convert.py

@@ -0,0 +1,35 @@
+"""
+A sample of convert the cifar100 dataset to 224 * 224 size train\val data.
+"""
+import cv2
+import os
+from keras.datasets import cifar100
+
+
+def convert():
+    train = 'train//'
+    val = 'validation//'
+
+    (X_train, y_train), (X_test, y_test) = cifar100.load_data(label_mode='fine')
+
+    for i in range(len(X_train)):
+        x = X_train[i]
+        y = y_train[i]
+        path = train + str(y[0])
+        x = cv2.resize(x, (224, 224), interpolation=cv2.INTER_CUBIC)
+        if not os.path.exists(path):
+            os.makedirs(path)
+        cv2.imwrite(path + '//' + str(i) + '.jpg', x)
+
+    for i in range(len(X_test)):
+        x = X_train[i]
+        y = y_train[i]
+        path = val + str(y[0])
+        x = cv2.resize(x, (224, 224), interpolation=cv2.INTER_CUBIC)
+        if not os.path.exists(path):
+            os.makedirs(path)
+        cv2.imwrite(path + '//' + str(i) + '.jpg', x)
+
+
+if __name__ == '__main__':
+    convert()

mobilenet_v2.py

@@ -8,12 +8,11 @@
 
 
 from keras.models import Model
-from keras.layers import Input, Conv2D, AveragePooling2D, Dropout
-from keras.layers import Activation, BatchNormalization, add
+from keras.layers import Input, Conv2D, GlobalAveragePooling2D, Dropout
+from keras.layers import Activation, BatchNormalization, add, Reshape
 from keras.applications.mobilenet import relu6, DepthwiseConv2D
 from keras.utils.vis_utils import plot_model
 
-
 from keras import backend as K
 
 
@@ -130,11 +129,13 @@ def MobileNetv2(input_shape, k):
     x = _inverted_residual_block(x, 320, (3, 3), t=6, strides=1, n=1)
 
     x = _conv_block(x, 1280, (1, 1), strides=(1, 1))
-    # x = GlobalAveragePooling2D()(x)
-    x = AveragePooling2D((int(x.shape[1]), int(x.shape[2])))(x)
-    x = Dropout(0.3)(x)
+    x = GlobalAveragePooling2D()(x)
+    x = Reshape((1, 1, 1280))(x)
+    x = Dropout(0.3, name='Dropout')(x)
     x = Conv2D(k, (1, 1), padding='same')(x)
-    output = Activation('softmax', name='softmax')(x)
+
+    x = Activation('softmax', name='softmax')(x)
+    output = Reshape((k,))(x)
 
     model = Model(inputs, output)
     plot_model(model, to_file='images/MobileNetv2.png', show_shapes=True)