Merge pull request #36 from Oneflow-Inc/add_variable_batch_size

add variable batch_size function

examples/oneflow2onnx/variable_batch_size/test_mobilenet_v2_batch_none.py

+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import oneflow as flow
+import oneflow.typing as tp
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+# ref : https://arxiv.org/pdf/1801.04381.pdf
+# ref : https://github.com/liangfu/mxnet-mobilenet-v2/blob/master/symbols/mobilenetv2.py
+
+
+def _get_regularizer(model_name):
+    # all decay
+    return flow.regularizers.l2(0.00004)
+
+
+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,
+    training=True,
+    name=None,
+):
+    return flow.layers.batch_normalization(
+        inputs=inputs,
+        axis=axis,
+        momentum=momentum,
+        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=training,
+        name=name,
+    )
+
+
+def _relu6(data, prefix):
+    return flow.clip_by_value(data, 0, 6, name="%s-relu6" % prefix)
+
+
+def mobilenet_unit(
+    data,
+    num_filter=1,
+    kernel=(1, 1),
+    stride=(1, 1),
+    pad=(0, 0),
+    num_group=1,
+    data_format="NCHW",
+    if_act=True,
+    use_bias=False,
+    trainable=True,
+    training=True,
+    prefix="",
+):
+    conv = flow.layers.conv2d(
+        inputs=data,
+        filters=num_filter,
+        kernel_size=kernel,
+        strides=stride,
+        padding=pad,
+        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,
+    )
+    bn = _batch_norm(
+        conv,
+        axis=1,
+        momentum=0.9,
+        epsilon=1e-5,
+        trainable=trainable,
+        training=training,
+        name="%s-BatchNorm" % prefix,
+    )
+    if if_act:
+        act = _relu6(bn, prefix)
+        return act
+    else:
+        return bn
+
+
+def shortcut(data_in, data_residual, prefix):
+    out = flow.math.add(data_in, data_residual)
+    return out
+
+
+def inverted_residual_unit(
+    data,
+    num_in_filter,
+    num_filter,
+    ifshortcut,
+    stride,
+    kernel,
+    pad,
+    expansion_factor,
+    prefix,
+    trainable=True,
+    training=True,
+    data_format="NCHW",
+    has_expand=1,
+):
+    num_expfilter = int(round(num_in_filter * expansion_factor))
+    if has_expand:
+        channel_expand = mobilenet_unit(
+            data=data,
+            num_filter=num_expfilter,
+            kernel=(1, 1),
+            stride=(1, 1),
+            pad="valid",
+            num_group=1,
+            data_format=data_format,
+            if_act=True,
+            trainable=trainable,
+            training=training,
+            prefix="%s-expand" % prefix,
+        )
+    else:
+        channel_expand = data
+    bottleneck_conv = mobilenet_unit(
+        data=channel_expand,
+        num_filter=num_expfilter,
+        stride=stride,
+        kernel=kernel,
+        pad=pad,
+        num_group=num_expfilter,
+        data_format=data_format,
+        if_act=True,
+        trainable=trainable,
+        training=training,
+        prefix="%s-depthwise" % prefix,
+    )
+    linear_out = mobilenet_unit(
+        data=bottleneck_conv,
+        num_filter=num_filter,
+        kernel=(1, 1),
+        stride=(1, 1),
+        pad="valid",
+        num_group=1,
+        data_format=data_format,
+        if_act=False,
+        trainable=trainable,
+        training=training,
+        prefix="%s-project" % prefix,
+    )
+
+    if ifshortcut:
+        out = shortcut(data_in=data, data_residual=linear_out, prefix=prefix,)
+        return out
+    else:
+        return linear_out
+
+
+MNETV2_CONFIGS_MAP = {
+    (224, 224): {
+        "firstconv_filter_num": 32,
+        # t, c, s
+        "bottleneck_params_list": [
+            (1, 16, 1, False),
+            (6, 24, 2, False),
+            (6, 24, 1, True),
+            (6, 32, 2, False),
+            (6, 32, 1, True),
+            (6, 32, 1, True),
+            (6, 64, 2, False),
+            (6, 64, 1, True),
+            (6, 64, 1, True),
+            (6, 64, 1, True),
+            (6, 96, 1, False),
+            (6, 96, 1, True),
+            (6, 96, 1, True),
+            (6, 160, 2, False),
+            (6, 160, 1, True),
+            (6, 160, 1, True),
+            (6, 320, 1, False),
+        ],
+        "filter_num_before_gp": 1280,
+    }
+}
+
+
+class MobileNetV2(object):
+    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:
+            self.config_map = MNETV2_CONFIGS_MAP[(224, 224)]
+
+    def build_network(
+        self, input_data, data_format, class_num=1000, prefix="", **configs
+    ):
+        self.config_map.update(configs)
+
+        first_c = int(round(self.config_map["firstconv_filter_num"] * self.multiplier))
+        first_layer = mobilenet_unit(
+            data=input_data,
+            num_filter=first_c,
+            kernel=(3, 3),
+            stride=(2, 2),
+            pad="same",
+            data_format=data_format,
+            if_act=True,
+            trainable=self.trainable,
+            training=self.training,
+            prefix=prefix + "-Conv",
+        )
+
+        last_bottleneck_layer = first_layer
+        in_c = first_c
+        for i, layer_setting in enumerate(self.config_map["bottleneck_params_list"]):
+            t, c, s, sc = layer_setting
+            if i == 0:
+                last_bottleneck_layer = inverted_residual_unit(
+                    data=last_bottleneck_layer,
+                    num_in_filter=in_c,
+                    num_filter=int(round(c * self.multiplier)),
+                    ifshortcut=sc,
+                    stride=(s, s),
+                    kernel=(3, 3),
+                    pad="same",
+                    expansion_factor=t,
+                    prefix=prefix + "-expanded_conv",
+                    trainable=self.trainable,
+                    training=self.training,
+                    data_format=data_format,
+                    has_expand=0,
+                )
+                in_c = int(round(c * self.multiplier))
+            else:
+                last_bottleneck_layer = inverted_residual_unit(
+                    data=last_bottleneck_layer,
+                    num_in_filter=in_c,
+                    num_filter=int(round(c * self.multiplier)),
+                    ifshortcut=sc,
+                    stride=(s, s),
+                    kernel=(3, 3),
+                    pad="same",
+                    expansion_factor=t,
+                    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,
+            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",
+        )
+        # 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="", **configs):
+        sym = self.build_network(
+            input_data, class_num=class_num, prefix=prefix, **configs
+        )
+        return sym
+
+
+def Mobilenet(
+    input_data,
+    channel_last=False,
+    trainable=True,
+    training=True,
+    num_classes=1000,
+    multiplier=1.0,
+    prefix="",
+):
+    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 out
+
+
+def test_mobilenetv2():
+    @flow.global_function()
+    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", dynamic_batch_size=True)
+

oneflow_onnx/oneflow2onnx/flow2onnx.py

@@ -230,6 +230,7 @@ def Export(
     extra_opset: Optional[int] = None,
     shape_override: Optional[Dict[Text, List[int]]] = None,
     external_data: bool = False,
+    dynamic_batch_size: bool = False,
 ):
     r"""Export a oneflow model into ONNX format.
 
@@ -262,6 +263,10 @@ def Export(
             model_proto = onnx_graph.MakeModel(
                 job_name, onnx_filename, external_data=external_data
             )
+
+            if dynamic_batch_size == True:
+                model_proto.graph.input[0].type.tensor_type.shape.dim[0].dim_param = 'None'
+
             with open(onnx_filename, "wb") as f:
                 try:
                     f.write(model_proto.SerializeToString())

oneflow_onnx/oneflow2onnx/util.py

@@ -64,6 +64,7 @@ def export_onnx_model(
     opset=None,
     flow_weight_dir=None,
     onnx_model_path="/tmp",
+    dynamic_batch_size=False,
 ):
     if flow_weight_dir == None:
         flow_weight_dir = tempfile.TemporaryDirectory()
@@ -79,6 +80,7 @@ def export_onnx_model(
             onnx_model_path,
             opset=opset,
             external_data=external_data,
+            dynamic_batch_size=dynamic_batch_size,
         )
         flow_weight_dir.cleanup()
     else:
@@ -92,6 +94,7 @@ def export_onnx_model(
             onnx_model_path,
             opset=opset,
             external_data=external_data,
+            dynamic_batch_size=dynamic_batch_size,
         )
 
     def cleanup():
@@ -126,6 +129,7 @@ def convert_to_onnx_and_check(
     opset=None,
     flow_weight_dir=None,
     onnx_model_path="/tmp",
+    dynamic_batch_size=False,
 ):
     if explicit_init:
         # it is a trick to keep check_point.save() from hanging when there is no variable
@@ -141,17 +145,20 @@ def convert_to_onnx_and_check(
     flow.train.CheckPoint().init()
 
     onnx_model_path, cleanup = export_onnx_model(
-        job_func, external_data, opset, flow_weight_dir, onnx_model_path
+        job_func, external_data, opset, flow_weight_dir, onnx_model_path, dynamic_batch_size
     )
 
-    ipt_dict, onnx_res = run_onnx(
+
+    if dynamic_batch_size != True:
+        ipt_dict, onnx_res = run_onnx(
         onnx_model_path, ["CPUExecutionProvider"], ort_optimize=ort_optimize
-    )
-    oneflow_res = job_func(*ipt_dict.values())
-    if not isinstance(oneflow_res, np.ndarray):
-        oneflow_res = oneflow_res.get().numpy()
+        )
+        oneflow_res = job_func(*ipt_dict.values())
+        if not isinstance(oneflow_res, np.ndarray):
+            oneflow_res = oneflow_res.get().numpy()
+
+        compare_result(oneflow_res, onnx_res, print_outlier=print_outlier)
 
-    compare_result(oneflow_res, onnx_res, print_outlier=print_outlier)
 
     flow.clear_default_session()
     # cleanup()