add ReflectionPad2d

docs/source/experimental.rst

@@ -202,6 +202,7 @@ Experimental features
 .. autofunction:: oneflow.experimental.Tensor.ceil
 .. autofunction:: oneflow.experimental.expm1
 .. autofunction:: oneflow.experimental.Tensor.expm1
+.. autofunction:: oneflow.experimental.nn.ReflectionPad2d
 .. autofunction:: oneflow.experimental.meshgrid
 .. autofunction:: oneflow.experimental.topk
-.. autofunction:: oneflow.experimental.Tensor.topk
+.. autofunction:: oneflow.experimental.Tensor.topk

oneflow/python/nn/modules/reflection_pad2d.py

@@ -0,0 +1,110 @@
+"""
+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
+from oneflow.python.oneflow_export import oneflow_export, experimental_api
+from oneflow.python.nn.module import Module
+from typing import Union
+
+
+@oneflow_export("nn.ReflectionPad2d")
+@experimental_api
+class ReflectionPad2d(Module):
+    r"""The interface is consistent with PyTorch.
+    The documentation is referenced from:
+    https://pytorch.org/docs/stable/generated/torch.nn.ReflectionPad2d.html
+
+
+    This operator pads the input tensor using the reflection of the input boundary.
+
+    Args:
+        padding (Union[int,tuple]): The size or bundary of padding, if is `int` uses the same padding in all dimension; if 4-dims `tuple`, uses :math:`(\text{padding}_{\text{left}}, \text{padding}_{\text{right}}, \text{padding}_{\text{top}}, \text{padding}_{\text{bottom}} )`
+
+    Returns:
+        Tensor: Returns a new tensor which is result of the reflection padding of the input tensor.
+
+    Shape:
+        - Input: :math:`(N, C, H_{\text{in}}, W_{\text{in}})`
+        - Output: :math:`(N, C, H_{\text{out}}, W_{\text{out}})` where
+
+          :math:`H_{\text{out}} = H_{\text{in}} + \text{padding}_{\text{top}} + \text{padding}_{\text{bottom}}`
+
+          :math:`W_{\text{out}} = W_{\text{in}} + \text{padding}_{\text{left}} + \text{padding}_{\text{right}}`
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow.experimental as flow
+        >>> import numpy as np
+        >>> flow.enable_eager_execution()
+        >>> input = flow.Tensor(np.arange(18).reshape((1, 2, 3, 3)), dtype=flow.float32)
+        >>> m = flow.nn.ReflectionPad2d((2, 2, 1, 1))
+        >>> out = m(input)
+        >>> out
+        tensor([[[[ 5.,  4.,  3.,  4.,  5.,  4.,  3.],
+                  [ 2.,  1.,  0.,  1.,  2.,  1.,  0.],
+                  [ 5.,  4.,  3.,  4.,  5.,  4.,  3.],
+                  [ 8.,  7.,  6.,  7.,  8.,  7.,  6.],
+                  [ 5.,  4.,  3.,  4.,  5.,  4.,  3.]],
+        <BLANKLINE>         
+                 [[14., 13., 12., 13., 14., 13., 12.],
+                  [11., 10.,  9., 10., 11., 10.,  9.],
+                  [14., 13., 12., 13., 14., 13., 12.],
+                  [17., 16., 15., 16., 17., 16., 15.],
+                  [14., 13., 12., 13., 14., 13., 12.]]]], dtype=oneflow.float32)
+
+    """
+
+    def __init__(self, padding: Union[int, tuple]) -> None:
+        super().__init__()
+        if isinstance(padding, tuple):
+            assert len(padding) == 4, ValueError("Padding length must be 4")
+            boundary = [padding[0], padding[1], padding[2], padding[3]]
+        elif isinstance(padding, int):
+            boundary = [padding, padding, padding, padding]
+        else:
+            raise ValueError("padding must be in or list or tuple!")
+        self.padding = boundary
+        self._op = (
+            flow.builtin_op("reflection_pad2d")
+            .Input("x")
+            .Output("y")
+            .Attr("padding", boundary)
+            .Attr("floating_value", float(1.0))
+            .Attr("integral_value", int(0))
+            .Build()
+        )
+
+    def forward(self, x):
+        H, W = x.shape[2], x.shape[3]
+        if (
+            self.padding[2] < H
+            and self.padding[3] < H
+            and self.padding[0] < W
+            and self.padding[1] < W
+        ):
+            res = self._op(x)[0]
+            return res
+        else:
+            raise ValueError(
+                "padding size should be less than the corresponding input dimension!"
+            )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(raise_on_error=True)

oneflow/python/ops/pad.py

@@ -20,7 +20,7 @@
 import oneflow
 import oneflow.python.framework.id_util as id_util
 import oneflow.python.framework.remote_blob as remote_blob_util
-from oneflow.python.oneflow_export import oneflow_export
+from oneflow.python.oneflow_export import oneflow_export, stable_api
 import oneflow._oneflow_internal
 
 
@@ -238,6 +238,7 @@ def same_pad_Job(x: tp.Numpy.Placeholder((1, 1, 3, 3))
 
 
 @oneflow_export("reflection_pad2d")
+@stable_api
 def reflection_pad2d(
     x: oneflow._oneflow_internal.BlobDesc,
     padding: Union[int, tuple, list],

oneflow/python/test/modules/test_reflection_pad2d.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 unittest
+from collections import OrderedDict
+import oneflow.experimental as flow
+import numpy as np
+
+from test_util import (
+    GenArgList,
+    FlattenArray,
+    Array2Numpy,
+    Index2Coordinate,
+)
+
+
+def gen_numpy_test_sample(input, padding):
+    c_idx, h_idx, w_idx = 1, 2, 3
+    pad_left = padding[0]
+    pad_right = padding[1]
+    pad_top = padding[2]
+    pad_bottom = padding[3]
+    pad_shape = ((0, 0), (0, 0), (pad_top, pad_bottom), (pad_left, pad_right))
+
+    def _np_reflection_pad2d(input, pad_shape):
+        numpy_reflect = np.pad(input, pad_shape, "reflect")
+        return numpy_reflect
+
+    def _np_reflection_pad2d_grad(src, dest):
+        dx_height, dx_width = input.shape[h_idx], input.shape[w_idx]
+        dy_height, dy_width = output.shape[h_idx], output.shape[w_idx]
+
+        numpy_src = np.ones(src.shape, np.int32)
+        numpy_dest = np.zeros(dest.shape, np.int32)
+        array_src = FlattenArray(numpy_src)
+        array_dest = FlattenArray(numpy_dest)
+
+        src_num = src.shape[c_idx] * src.shape[h_idx] * src.shape[w_idx]
+        dest_num = dest.shape[c_idx] * dest.shape[h_idx] * dest.shape[w_idx]
+        elements_num = src.shape[0] * src_num
+        for iter_n in range(elements_num):
+            coords = Index2Coordinate(iter_n, src.shape)
+            n, c, i, j = coords[0], coords[c_idx], coords[h_idx], coords[w_idx]
+            ip_x = ip_y = 0
+            if j < pad_left:
+                ip_x = pad_left * 2 - j
+            elif j >= pad_left and j < (dx_width + pad_left):
+                ip_x = j
+            else:
+                ip_x = (dx_width + pad_left - 1) * 2 - j
+
+            if i < pad_top:
+                ip_y = pad_top * 2 - i
+            elif i >= pad_top and i < (dx_height + pad_top):
+                ip_y = i
+            else:
+                ip_y = (dx_height + pad_top - 1) * 2 - i
+
+            ip_x = ip_x - pad_left
+            ip_y = ip_y - pad_top
+            src_index = n * src_num + c * dy_width * dy_height + i * dy_width + j
+            dest_index = (
+                n * dest_num + c * dx_width * dx_height + ip_y * dx_width + ip_x
+            )
+            array_dest[dest_index] += array_src[src_index]
+
+        numpy_dest = Array2Numpy(array_dest, dest.shape)
+        return numpy_dest
+
+    output = _np_reflection_pad2d(input, pad_shape)
+    grad = _np_reflection_pad2d_grad(output, input)
+
+    return output, grad
+
+
+def _test_reflection_pad2d(test_case, shape, padding, device):
+    np_input = np.random.randn(*shape).astype(np.float32)
+
+    of_input = flow.Tensor(
+        np_input, dtype=flow.float32, device=flow.device(device), requires_grad=True
+    )
+
+    if isinstance(padding, int):
+        boundary = [padding, padding, padding, padding]
+
+    elif isinstance(padding, tuple) and len(padding) == 4:
+        boundary = [padding[0], padding[1], padding[2], padding[3]]
+    else:
+        raise ValueError("padding must be in or list or tuple!")
+    np_out, np_grad = gen_numpy_test_sample(np_input, boundary)
+
+    layer = flow.nn.ReflectionPad2d(padding=padding)
+    of_out = layer(of_input)
+    test_case.assertTrue(np.allclose(of_out.numpy(), np_out, 1e-4, 1e-4))
+    of_out = of_out.sum()
+    of_out.backward()
+    test_case.assertTrue(np.allclose(of_input.grad.numpy(), np_grad, 1e-4, 1e-4))
+
+
+@unittest.skipIf(
+    not flow.unittest.env.eager_execution_enabled(),
+    ".numpy() doesn't work in lazy mode",
+)
+class TestReflectionPad2dModule(flow.unittest.TestCase):
+    def test_reflection_pad2d(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["shape"] = [(1, 2, 3, 4), (8, 3, 4, 4)]
+        arg_dict["padding"] = [(2), (1, 1, 2, 2)]
+        arg_dict["device"] = ["cpu", "cuda"]
+        for arg in GenArgList(arg_dict):
+            _test_reflection_pad2d(test_case, *arg)
+
+
+if __name__ == "__main__":
+    unittest.main()