add mish activation

oneflow/python/ops/math_ops.py

@@ -2105,3 +2105,47 @@ def range_job()->tp.Numpy:
         .InferAndTryRun()
         .RemoteBlobList()[0]
     )
+
+
+@oneflow_export("math.mish")
+def mish(
+    x: remote_blob_util.BlobDef, name: Optional[str] = None,
+) -> remote_blob_util.BlobDef:
+    """The Mish activation function. 
+
+    The equation is: 
+
+    .. math:: 
+
+        out = x*tanh(ln(1+e^x))
+
+    For example: 
+
+    .. code-block:: python 
+
+        import oneflow as flow
+        import oneflow.typing as tp 
+        import numpy as np 
+
+
+        @flow.global_function()
+        def mish_job(x: tp.Numpy.Placeholder(shape=(5, )))->tp.Numpy: 
+            return flow.math.mish(x)
+
+
+        x = np.array([-0.5, 0, 0.5, 1.0, 1.5]).astype(np.float32)
+        out = mish_job(x)
+
+    Args:
+        x (remote_blob_util.BlobDef): The input Blob. 
+        name (Optional[str], optional): The name for the operation. Defaults to None.
+
+    Returns:
+        remote_blob_util.BlobDef: The result Blob. 
+    """
+    if name is None:
+        name = id_util.UniqueStr("Mish_")
+
+    return x * flow.math.tanh(
+        flow.math.softplus(x, name=name + "softplus"), name=name + "tanh"
+    )

oneflow/python/test/ops/test_mish.py

@@ -0,0 +1,127 @@
+"""
+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 numpy as np
+import oneflow.typing as tp
+from test_util import GenArgList
+import unittest
+from collections import OrderedDict
+from typing import Dict
+import os
+
+
+def _compare_mish_with_np(input_shape, device_type, machine_ids, device_counts):
+    input_1 = np.random.random(size=input_shape).astype(np.float32)
+
+    assert device_type in ["cpu", "gpu"]
+
+    flow.clear_default_session()
+    if device_type == "cpu":
+        flow.config.cpu_device_num(device_counts)
+    else:
+        flow.config.gpu_device_num(device_counts)
+
+    func_config = flow.FunctionConfig()
+
+    def np_mish(input):
+        return input * np.tanh(np.log1p(np.exp(input)))
+
+    np_out_mish = np_mish(input_1)
+
+    def np_diff(input):
+        u = np.log1p(np.exp(input))
+
+        return np.tanh(u) + input * (1 - np.tanh(u) ** 2) * (
+            np.exp(input) / (1 + np.exp(input))
+        )
+
+    _np_grad = np_diff(input_1)
+
+    def assert_prediction_grad(blob: tp.Numpy):
+        assert np.allclose(blob, _np_grad)
+
+    @flow.global_function(
+        type="train", function_config=func_config,
+    )
+    def oneflow_mish(
+        of_input_1: tp.Numpy.Placeholder(shape=input_1.shape),
+    ) -> tp.Numpy:
+        with flow.scope.placement(device_type, "0:0"):
+            v = flow.get_variable(
+                shape=input_1.shape,
+                dtype=flow.float32,
+                initializer=flow.zeros_initializer(),
+                name="x_var",
+            )
+            x_var = of_input_1 + v
+
+        flow.watch_diff(x_var, assert_prediction_grad)
+
+        of_mish_out = flow.math.mish(x_var)
+
+        with flow.scope.placement(device_type, "0:0"):
+            flow.optimizer.SGD(
+                flow.optimizer.PiecewiseConstantScheduler([], [1e-3]), momentum=0
+            ).minimize(of_mish_out)
+
+        return of_mish_out
+
+    of_out_mish = oneflow_mish(input_1)
+
+    assert np.allclose(of_out_mish, np_out_mish)
+
+
+def _gen_arg_dict(shape, device_type, machine_ids, device_counts):
+    # Generate a dict to pass parameter to test case
+    arg_dict = OrderedDict()
+    arg_dict["input_shape"] = [shape]
+    arg_dict["device_type"] = [device_type]
+    arg_dict["machine_ids"] = [machine_ids]
+    arg_dict["device_counts"] = [device_counts]
+    return arg_dict
+
+
+@flow.unittest.skip_unless_1n1d()
+class Testmish1n1d(flow.unittest.TestCase):
+    def test_mish_cpu(test_case):
+        arg_dict = _gen_arg_dict(
+            shape=(3, 3), device_type="cpu", machine_ids="0:0", device_counts=1
+        )
+        for arg in GenArgList(arg_dict):
+            _compare_mish_with_np(*arg)
+
+    @unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")
+    def test_mish_gpu(test_case):
+        arg_dict = _gen_arg_dict(
+            shape=(3, 16, 32), device_type="gpu", machine_ids="0:0", device_counts=1,
+        )
+        for arg in GenArgList(arg_dict):
+            _compare_mish_with_np(*arg)
+
+
+@flow.unittest.skip_unless_1n2d()
+class Testmish1n2d(flow.unittest.TestCase):
+    @unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases")
+    def test_mish_gpu_1n2d(test_case):
+        arg_dict = _gen_arg_dict(
+            shape=(3, 8, 8, 4), device_type="gpu", machine_ids="0:0-1", device_counts=2,
+        )
+        for arg in GenArgList(arg_dict):
+            _compare_mish_with_np(*arg)
+
+
+if __name__ == "__main__":
+    unittest.main()