Merge pull request #33 from mosout/fix_mobilenet_v2

fix mobilenet v2

examples/oneflow2onnx/models/test_mobilenet_v2.py

@@ -15,83 +15,43 @@ limitations under the License.
 """
 import oneflow as flow
 import oneflow.typing as tp
-import oneflow.core.operator.op_conf_pb2 as op_conf_util
 from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
 
-g_trainable = False
+# ref : https://arxiv.org/pdf/1801.04381.pdf
+# ref : https://github.com/liangfu/mxnet-mobilenet-v2/blob/master/symbols/mobilenetv2.py
 
 
-def _conv2d_layer(
-    name,
-    input,
-    filters,
-    kernel_size=3,
-    strides=1,
-    padding="SAME",
-    group_num=1,
-    data_format="NCHW",
-    dilation_rate=1,
-    activation=op_conf_util.kRelu,
-    use_bias=False,
-    weight_initializer=flow.random_uniform_initializer(),
-    bias_initializer=flow.random_uniform_initializer(),
-):
+def _get_regularizer(model_name):
+    # all decay
+    return flow.regularizers.l2(0.00004)
 
-    if data_format == "NCHW":
-        weight_shape = (
-            int(filters),
-            int(input.shape[1] / group_num),
-            int(kernel_size[0]),
-            int(kernel_size[0]),
-        )
-    elif data_format == "NHWC":
-        weight_shape = (
-            int(filters),
-            int(kernel_size[0]),
-            int(kernel_size[0]),
-            int(input.shape[3] / group_num),
-        )
-    else:
-        raise ValueError('data_format must be "NCHW" or "NHWC".')
-    weight = flow.get_variable(
-        name + "-weight",
-        shape=weight_shape,
-        dtype=input.dtype,
-        initializer=weight_initializer,
-    )
-    output = flow.nn.conv2d(
-        input,
-        weight,
-        strides,
-        padding,
-        data_format,
-        dilation_rate,
-        group_num,
-        name=name,
-    )
-    if use_bias:
-        bias = flow.get_variable(
-            name + "-bias",
-            shape=(filters,),
-            dtype=input.dtype,
-            initializer=bias_initializer,
-            model_name="bias",
-        )
-        output = flow.nn.bias_add(output, bias, data_format)
-
-    if activation is not None:
-        if activation == op_conf_util.kRelu:
-            output = flow.keras.activations.relu(output)
-        else:
-            raise NotImplementedError
 
-    return output
+def _get_initializer(model_name):
+    if model_name == "weight":
+        return flow.variance_scaling_initializer(
+            2.0, mode="fan_out", distribution="random_normal", data_format="NCHW"
+        )
+    elif model_name == "bias":
+        return flow.zeros_initializer()
+    elif model_name == "gamma":
+        return flow.ones_initializer()
+    elif model_name == "beta":
+        return flow.zeros_initializer()
+    elif model_name == "dense_weight":
+        return flow.random_normal_initializer(0, 0.01)
 
 
 def _batch_norm(
-    inputs, axis, momentum, epsilon, center=True, scale=True, trainable=True, name=None
+    inputs,
+    axis,
+    momentum,
+    epsilon,
+    center=True,
+    scale=True,
+    trainable=True,
+    training=True,
+    name=None,
 ):
-
     return flow.layers.batch_normalization(
         inputs=inputs,
         axis=axis,
@@ -99,8 +59,12 @@ def _batch_norm(
         epsilon=epsilon,
         center=center,
         scale=scale,
+        beta_initializer=_get_initializer("beta"),
+        gamma_initializer=_get_initializer("gamma"),
+        beta_regularizer=_get_regularizer("beta"),
+        gamma_regularizer=_get_regularizer("gamma"),
         trainable=trainable,
-        training=trainable,
+        training=training,
         name=name,
     )
 
@@ -119,35 +83,35 @@ def mobilenet_unit(
     data_format="NCHW",
     if_act=True,
     use_bias=False,
+    trainable=True,
+    training=True,
     prefix="",
 ):
-    conv = _conv2d_layer(
-        name=prefix,
-        input=data,
+    conv = flow.layers.conv2d(
+        inputs=data,
         filters=num_filter,
         kernel_size=kernel,
         strides=stride,
         padding=pad,
-        group_num=num_group,
         data_format=data_format,
         dilation_rate=1,
+        groups=num_group,
         activation=None,
         use_bias=use_bias,
+        kernel_initializer=_get_initializer("weight"),
+        bias_initializer=_get_initializer("bias"),
+        kernel_regularizer=_get_regularizer("weight"),
+        bias_regularizer=_get_regularizer("bias"),
+        name=prefix,
     )
-    if data_format == "NCHW":
-        axis = 1
-    elif data_format == "NHWC":
-        axis = 3
-    else:
-        raise ValueError('data_format must be "NCHW" or "NHWC".')
-
     bn = _batch_norm(
         conv,
-        axis=axis,
-        momentum=0.97,
-        epsilon=1e-3,
+        axis=1,
+        momentum=0.9,
+        epsilon=1e-5,
+        trainable=trainable,
+        training=training,
         name="%s-BatchNorm" % prefix,
-        trainable=g_trainable,
     )
     if if_act:
         act = _relu6(bn, prefix)
@@ -156,33 +120,6 @@ def mobilenet_unit(
         return bn
 
 
-def conv(
-    data,
-    num_filter=1,
-    kernel=(1, 1),
-    stride=(1, 1),
-    pad=(0, 0),
-    num_group=1,
-    data_format="NCHW",
-    use_bias=False,
-    prefix="",
-):
-    # return _conv2d_layer(name='%s-conv2d'%prefix, input=data, filters=num_filter, kernel_size=kernel, strides=stride, padding=pad, group_num=num_group, dilation_rate=1, activation=None, use_bias=use_bias)
-    return _conv2d_layer(
-        name=prefix,
-        input=data,
-        filters=num_filter,
-        kernel_size=kernel,
-        strides=stride,
-        padding=pad,
-        group_num=num_group,
-        data_format=data_format,
-        dilation_rate=1,
-        activation=None,
-        use_bias=use_bias,
-    )
-
-
 def shortcut(data_in, data_residual, prefix):
     out = flow.math.add(data_in, data_residual)
     return out
@@ -198,6 +135,8 @@ def inverted_residual_unit(
     pad,
     expansion_factor,
     prefix,
+    trainable=True,
+    training=True,
     data_format="NCHW",
     has_expand=1,
 ):
@@ -212,6 +151,8 @@ def inverted_residual_unit(
             num_group=1,
             data_format=data_format,
             if_act=True,
+            trainable=trainable,
+            training=training,
             prefix="%s-expand" % prefix,
         )
     else:
@@ -225,6 +166,8 @@ def inverted_residual_unit(
         num_group=num_expfilter,
         data_format=data_format,
         if_act=True,
+        trainable=trainable,
+        training=training,
         prefix="%s-depthwise" % prefix,
     )
     linear_out = mobilenet_unit(
@@ -236,6 +179,8 @@ def inverted_residual_unit(
         num_group=1,
         data_format=data_format,
         if_act=False,
+        trainable=trainable,
+        training=training,
         prefix="%s-project" % prefix,
     )
 
@@ -275,10 +220,12 @@ MNETV2_CONFIGS_MAP = {
 
 
 class MobileNetV2(object):
-    def __init__(self, data_wh, multiplier, **kargs):
+    def __init__(self, data_wh, multiplier, trainable=True, training=True, **kargs):
         super(MobileNetV2, self).__init__()
         self.data_wh = data_wh
         self.multiplier = multiplier
+        self.trainable = trainable
+        self.training = training
         if self.data_wh in MNETV2_CONFIGS_MAP:
             self.config_map = MNETV2_CONFIGS_MAP[self.data_wh]
         else:
@@ -289,11 +236,6 @@ class MobileNetV2(object):
     ):
         self.config_map.update(configs)
 
-        # input_data = flow.math.multiply(input_data, 2.0 / 255.0)
-        # input_data = flow.math.add(input_data, -1)
-
-        if data_format == "NCHW":
-            input_data = flow.transpose(input_data, name="transpose", perm=[0, 3, 1, 2])
         first_c = int(round(self.config_map["firstconv_filter_num"] * self.multiplier))
         first_layer = mobilenet_unit(
             data=input_data,
@@ -303,6 +245,8 @@ class MobileNetV2(object):
             pad="same",
             data_format=data_format,
             if_act=True,
+            trainable=self.trainable,
+            training=self.training,
             prefix=prefix + "-Conv",
         )
 
@@ -321,6 +265,8 @@ class MobileNetV2(object):
                     pad="same",
                     expansion_factor=t,
                     prefix=prefix + "-expanded_conv",
+                    trainable=self.trainable,
+                    training=self.training,
                     data_format=data_format,
                     has_expand=0,
                 )
@@ -335,66 +281,79 @@ class MobileNetV2(object):
                     kernel=(3, 3),
                     pad="same",
                     expansion_factor=t,
-                    data_format=data_format,
                     prefix=prefix + "-expanded_conv_%d" % i,
+                    trainable=self.trainable,
+                    training=self.training,
+                    data_format=data_format,
                 )
                 in_c = int(round(c * self.multiplier))
-
         last_fm = mobilenet_unit(
             data=last_bottleneck_layer,
-            # num_filter=int(1280 * self.multiplier) if self.multiplier > 1.0 else 1280,
-            # gr to confirm
-            num_filter=int(256 * self.multiplier) if self.multiplier > 1.0 else 256,
+            num_filter=int(1280 * self.multiplier) if self.multiplier > 1.0 else 1280,
             kernel=(1, 1),
             stride=(1, 1),
             pad="valid",
             data_format=data_format,
             if_act=True,
+            trainable=self.trainable,
+            training=self.training,
             prefix=prefix + "-Conv_1",
         )
-        base_only = True
-        if base_only:
-            return last_fm
-        else:
-            raise NotImplementedError
+        # global average pooling
+        pool_size = int(self.data_wh[0] / 32)
+        pool = flow.nn.avg_pool2d(
+            last_fm,
+            ksize=pool_size,
+            strides=1,
+            padding="VALID",
+            data_format="NCHW",
+            name="pool5",
+        )
+        fc = flow.layers.dense(
+            flow.reshape(pool, (pool.shape[0], -1)),
+            units=class_num,
+            use_bias=False,
+            kernel_initializer=_get_initializer("dense_weight"),
+            bias_initializer=_get_initializer("bias"),
+            kernel_regularizer=_get_regularizer("dense_weight"),
+            bias_regularizer=_get_regularizer("bias"),
+            trainable=self.trainable,
+            name=prefix + "-fc",
+        )
+        return fc
 
-    def __call__(
-        self, input_data, class_num=1000, prefix="", layer_out=None, **configs
-    ):
+    def __call__(self, input_data, class_num=1000, prefix="", **configs):
         sym = self.build_network(
             input_data, class_num=class_num, prefix=prefix, **configs
         )
-        if layer_out is None:
-            return sym
-
-        internals = sym.get_internals()
-        if type(layer_out) is list or type(layer_out) is tuple:
-            layers_out = [
-                internals[layer_nm.strip() + "_output"] for layer_nm in layer_out
-            ]
-            return layers_out
-        else:
-            layer_out = internals[layer_out.strip() + "_output"]
-            return layer_out
+        return sym
 
 
 def Mobilenet(
-    input_data, data_format="NCHW", num_classes=1000, multiplier=1.0, prefix=""
+    input_data,
+    channel_last=False,
+    trainable=True,
+    training=True,
+    num_classes=1000,
+    multiplier=1.0,
+    prefix="",
 ):
-    mobilenetgen = MobileNetV2((224, 224), multiplier=multiplier)
-    layer_out = mobilenetgen(
-        input_data,
-        data_format=data_format,
-        class_num=num_classes,
-        prefix=prefix + "-MobilenetV2",
-        layer_out=None,
+    assert (
+        channel_last == False
+    ), "Mobilenet does not support channel_last mode, set channel_last=False will be right!"
+    data_format = "NCHW"
+    mobilenetgen = MobileNetV2(
+        (224, 224), multiplier=multiplier, trainable=trainable, training=training
+    )
+    out = mobilenetgen(
+        input_data, data_format=data_format, class_num=num_classes, prefix="MobilenetV2"
     )
-    return layer_out
+    return out
 
 
 def test_mobilenetv2():
     @flow.global_function()
-    def mobilenetv2(x: tp.Numpy.Placeholder((1, 224, 224, 3))):
+    def mobilenetv2(x: tp.Numpy.Placeholder((1, 3, 224, 224))):
         return Mobilenet(x)
 
     convert_to_onnx_and_check(mobilenetv2, flow_weight_dir=None, onnx_model_path="/tmp")