train mobilenetv1 for NHWC

zhangting2020 · Nov 6, 2019 · 4f0b3dd · 4f0b3dd
1 parent 9a10a36
commit 4f0b3dd
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 20 deletions.
diff --git a/PaddleCV/image_classification/models/mobilenet_v1.py b/PaddleCV/image_classification/models/mobilenet_v1.py
@@ -40,7 +40,8 @@ def net(self, input, class_dim=1000):
             num_filters=int(32 * scale),
             stride=2,
             padding=1,
-            name="conv1")
+            name="conv1",
+            data_format='NHWC')
 
         # 56x56
         input = self.depthwise_separable(
@@ -50,7 +51,8 @@ def net(self, input, class_dim=1000):
             num_groups=32,
             stride=1,
             scale=scale,
-            name="conv2_1")
+            name="conv2_1",
+            data_format='NHWC')
 
         input = self.depthwise_separable(
             input,
@@ -59,7 +61,8 @@ def net(self, input, class_dim=1000):
             num_groups=64,
             stride=2,
             scale=scale,
-            name="conv2_2")
+            name="conv2_2",
+            data_format='NHWC')
 
         # 28x28
         input = self.depthwise_separable(
@@ -69,7 +72,8 @@ def net(self, input, class_dim=1000):
             num_groups=128,
             stride=1,
             scale=scale,
-            name="conv3_1")
+            name="conv3_1",
+            data_format='NHWC')
 
         input = self.depthwise_separable(
             input,
@@ -78,7 +82,8 @@ def net(self, input, class_dim=1000):
             num_groups=128,
             stride=2,
             scale=scale,
-            name="conv3_2")
+            name="conv3_2",
+            data_format='NHWC')
 
         # 14x14
         input = self.depthwise_separable(
@@ -88,7 +93,8 @@ def net(self, input, class_dim=1000):
             num_groups=256,
             stride=1,
             scale=scale,
-            name="conv4_1")
+            name="conv4_1",
+            data_format='NHWC')
 
         input = self.depthwise_separable(
             input,
@@ -97,7 +103,8 @@ def net(self, input, class_dim=1000):
             num_groups=256,
             stride=2,
             scale=scale,
-            name="conv4_2")
+            name="conv4_2",
+            data_format='NHWC')
 
         # 14x14
         for i in range(5):
@@ -108,7 +115,8 @@ def net(self, input, class_dim=1000):
                 num_groups=512,
                 stride=1,
                 scale=scale,
-                name="conv5" + "_" + str(i + 1))
+                name="conv5" + "_" + str(i + 1),
+                data_format='NHWC')
         # 7x7
         input = self.depthwise_separable(
             input,
@@ -117,7 +125,8 @@ def net(self, input, class_dim=1000):
             num_groups=512,
             stride=2,
             scale=scale,
-            name="conv5_6")
+            name="conv5_6",
+            data_format='NHWC')
 
         input = self.depthwise_separable(
             input,
@@ -126,10 +135,11 @@ def net(self, input, class_dim=1000):
             num_groups=1024,
             stride=1,
             scale=scale,
-            name="conv6")
+            name="conv6",
+            data_format='NHWC')
 
         input = fluid.layers.pool2d(
-            input=input, pool_type='avg', global_pooling=True)
+            input=input, pool_type='avg', global_pooling=True, data_format='NHWC')
 
         output = fluid.layers.fc(input=input,
                                  size=class_dim,
@@ -148,7 +158,8 @@ def conv_bn_layer(self,
                       num_groups=1,
                       act='relu',
                       use_cudnn=True,
-                      name=None):
+                      name=None,
+                      data_format='NHWC'):
         conv = fluid.layers.conv2d(
             input=input,
             num_filters=num_filters,
@@ -160,15 +171,17 @@ def conv_bn_layer(self,
             use_cudnn=use_cudnn,
             param_attr=ParamAttr(
                 initializer=MSRA(), name=name + "_weights"),
-            bias_attr=False)
+            bias_attr=False,
+            data_format=data_format)
         bn_name = name + "_bn"
         return fluid.layers.batch_norm(
             input=conv,
             act=act,
             param_attr=ParamAttr(name=bn_name + "_scale"),
             bias_attr=ParamAttr(name=bn_name + "_offset"),
             moving_mean_name=bn_name + '_mean',
-            moving_variance_name=bn_name + '_variance')
+            moving_variance_name=bn_name + '_variance',
+            data_layout=data_format)
 
     def depthwise_separable(self,
                             input,
@@ -177,7 +190,8 @@ def depthwise_separable(self,
                             num_groups,
                             stride,
                             scale,
-                            name=None):
+                            name=None,
+                            data_format='NHWC'):
         depthwise_conv = self.conv_bn_layer(
             input=input,
             filter_size=3,
@@ -186,15 +200,17 @@ def depthwise_separable(self,
             padding=1,
             num_groups=int(num_groups * scale),
             use_cudnn=False,
-            name=name + "_dw")
+            name=name + "_dw",
+            data_format=data_format)
 
         pointwise_conv = self.conv_bn_layer(
             input=depthwise_conv,
             filter_size=1,
             num_filters=int(num_filters2 * scale),
             stride=1,
             padding=0,
-            name=name + "_sep")
+            name=name + "_sep",
+            data_format=data_format)
         return pointwise_conv
 
 

diff --git a/PaddleCV/image_classification/reader.py b/PaddleCV/image_classification/reader.py
@@ -228,9 +228,13 @@ def process_image(sample, settings, mode, color_jitter, rotate):
         img = policy(img)
         img = np.asarray(img)
 
-    img = img.astype('float32').transpose((2, 0, 1)) / 255
-    img_mean = np.array(mean).reshape((3, 1, 1))
-    img_std = np.array(std).reshape((3, 1, 1))
+    # img = img.astype('float32').transpose((2, 0, 1)) / 255
+    # img_mean = np.array(mean).reshape((3, 1, 1))
+    # img_std = np.array(std).reshape((3, 1, 1))
+    img = img.astype('float32') / 255
+    img_mean = np.array(mean).reshape((1, 1, 3))
+    img_std = np.array(std).reshape((1, 1, 3))
+
     img -= img_mean
     img /= img_std
 

diff --git a/PaddleCV/image_classification/train.sh b/PaddleCV/image_classification/train.sh
@@ -0,0 +1,10 @@
+python train.py \
+       --model=MobileNetV1 \
+       --batch_size=256 \
+       --total_images=1281167 \
+       --class_dim=1000 \
+       --image_shape=224,224,3 \
+       --model_save_dir=output/ \
+       --data_dir=/workspace/data/ \
+       --lr_strategy=piecewise_decay \
+       --lr=0.1