Add in_top_k function (#5428)

* Add in_top_k function

* auto format by CI

* fix bug

* Update oneflow/python/ops/math_ops.py

Co-authored-by: Yao Chi <later@usopp.net>

Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org>
Co-authored-by: Yao Chi <later@usopp.net>
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

docs/source/experimental.rst

@@ -240,3 +240,5 @@ Experimental features
 .. autofunction:: oneflow.experimental.Tensor.type_as
 .. autofunction:: oneflow.experimental.Tensor.long
 .. autofunction:: oneflow.experimental.bernoulli
+.. autofunction:: oneflow.experimental.in_top_k
+.. autofunction:: oneflow.experimental.Tensor.in_top_k
\ No newline at end of file

oneflow/python/nn/modules/in_top_k.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
+from oneflow.python.oneflow_export import oneflow_export, experimental_api
+from oneflow.python.framework.tensor import register_tensor_op
+from oneflow.python.nn.module import Module
+
+
+class InTopk(Module):
+    def __init__(self, k) -> None:
+        super().__init__()
+        self._in_top_k = (
+            flow.builtin_op("in_top_k")
+            .Input("targets")
+            .Input("predictions")
+            .Output("out")
+            .Attr("k", k)
+            .Build()
+        )
+
+    def forward(self, targets, predictions):
+        assert (
+            targets.shape[0] == predictions.shape[0]
+        ), "The num of targets must equal the num of predictions"
+        assert len(targets.shape) == 1, "The dimension of targets must be 1"
+        assert len(predictions.shape) == 2, "The dimension of predictions must be 2"
+        return self._in_top_k(targets, predictions)
+
+
+@oneflow_export("in_top_k")
+@experimental_api
+def in_top_k_op(targets, predictions, k):
+    r"""Says whether the targets are in the top K predictions.
+
+    Args:
+        targets (Tensor): the target tensor of type int32 or int64.
+        predictions (Tensor): the predictions tensor of type float32 .
+        k (int): Number of top elements to look at for computing precision.
+
+    Returns:
+        oneflow.Tensor: A Tensor of type bool. Computed Precision at k as a bool Tensor.
+
+    For example:
+
+    .. code-block:: python
+
+        >>> import oneflow.experimental as flow
+        >>> import numpy as np
+        >>> flow.enable_eager_execution()
+        >>> targets1 = flow.Tensor(np.array([3, 1]), dtype=flow.int32)
+        >>> predictions1 = flow.Tensor(np.array([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0],]), dtype=flow.float32)
+        >>> out1 = flow.in_top_k(targets1, predictions1, k=1)
+        >>> out1
+        tensor([1, 0], dtype=oneflow.int8)
+        >>> out2 = flow.in_top_k(targets1, predictions1, k=2)
+        >>> out2
+        tensor([1, 1], dtype=oneflow.int8)
+        >>> targets2 = flow.Tensor(np.array([3, 1]), dtype=flow.int32, device=flow.device('cuda'))
+        >>> predictions2 = flow.Tensor(np.array([[0.0, 1.0, 2.0, 3.0], [3.0, 2.0, 1.0, 0.0],]), dtype=flow.float32, device=flow.device('cuda'))
+        >>> out3 = flow.in_top_k(targets2, predictions2, k=1)
+        >>> out3
+        tensor([1, 0], device='cuda:0', dtype=oneflow.int8)
+
+    """
+
+    return InTopk(k=k)(targets, predictions)[0]
+
+
+@register_tensor_op("in_top_k")
+@experimental_api
+def in_top_k_op_tensor(targets, predictions, k):
+    r"""
+
+    in_top_k() -> Tensor
+
+    See :func:`oneflow.experimental.in_top_k`
+
+    """
+
+    return InTopk(k=k)(targets, predictions)[0]
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod(raise_on_error=True)

oneflow/python/ops/math_ops.py

@@ -2049,6 +2049,7 @@ def polyval(
 
 
 @oneflow_export("math.in_top_k", "in_top_k")
+@stable_api
 def in_top_k(
     targets: oneflow._oneflow_internal.BlobDesc,
     predictions: oneflow._oneflow_internal.BlobDesc,

oneflow/python/test/modules/test_in_top_k.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 unittest
+from collections import OrderedDict
+
+import numpy as np
+
+import oneflow.experimental as flow
+from test_util import GenArgList
+
+
+def _topk_np(input, k, dim: int = -1, largest: bool = True, _sorted: bool = True):
+    in_dims = input.shape
+    out_dims = list(in_dims)
+    num_axes = len(input.shape)
+    if dim < 0:
+        dim = dim + num_axes
+    n = in_dims[dim]
+    if k > n:
+        k = n
+    out_dims[dim] = k
+    out_dims = tuple(out_dims)
+    prev_dims = 1
+    next_dims = 1
+    for i in range(dim):
+        prev_dims *= in_dims[i]
+    for i in range(dim + 1, len(in_dims)):
+        next_dims *= in_dims[i]
+    input_flat = input.reshape((prev_dims, n, next_dims))
+
+    values_ref = np.ndarray(shape=(prev_dims, k, next_dims), dtype=input.dtype)
+    values_ref.fill(0)
+    indices_ref = np.ndarray(shape=(prev_dims, k, next_dims), dtype=np.int64)
+    indices_ref.fill(-1)
+    for i in range(prev_dims):
+        for j in range(next_dims):
+            kv = []
+            for x in range(n):
+                val = input_flat[i, x, j]
+                y = x * next_dims + i * in_dims[dim] * next_dims + j
+                kv.append((val, x, y))
+            cnt = 0
+            for val, x, y in sorted(kv, key=lambda x: (x[0], -x[1]), reverse=largest):
+                values_ref[i, cnt, j] = val
+                indices_ref[i, cnt, j] = x
+                cnt += 1
+                if cnt >= k or cnt >= n:
+                    break
+
+    values_ref = values_ref.reshape(out_dims)
+    indices_ref = indices_ref.reshape(out_dims)
+
+    return (values_ref, indices_ref)
+
+
+def _in_top_k_np(targets, predictions, k):
+    assert (
+        targets.shape[0] == predictions.shape[0]
+    ), "The num of targets must equal the num of predictions"
+    assert len(targets.shape) == 1, "The dimension of targets must be 1"
+    assert len(predictions.shape) == 2, "The dimension of predictions must be 2"
+    results = np.zeros_like(targets, dtype=np.int8)
+    for i in range(len(results)):
+        _, indices_topk = _topk_np(predictions[i], k)
+        if targets[i] in indices_topk:
+            results[i] = 1
+    return results
+
+
+def _test_in_top_k_impl(test_case, shape, k, device):
+    np_targets = np.random.randint(0, shape[1], size=shape[0])
+    np_predictions = np.random.rand(*shape)
+    of_targets = flow.Tensor(
+        np_targets, dtype=flow.int32, device=flow.device(device), requires_grad=True
+    )
+    of_predictions = flow.Tensor(
+        np_predictions,
+        dtype=flow.float32,
+        device=flow.device(device),
+        requires_grad=True,
+    )
+    of_out = flow.in_top_k(of_targets, of_predictions, k)
+    np_out = _in_top_k_np(np_targets, np_predictions, k)
+    test_case.assertTrue(
+        np.allclose(of_out.numpy(), np_out, 1e-4, 1e-4, equal_nan=True)
+    )
+
+
+@unittest.skipIf(
+    not flow.unittest.env.eager_execution_enabled(),
+    ".numpy() doesn't work in lazy mode",
+)
+class TestInTopK(flow.unittest.TestCase):
+    def test_in_top_k(test_case):
+        arg_dict = OrderedDict()
+        arg_dict["shape"] = [(2, 3), (3, 4), (5, 6)]
+        arg_dict["k"] = [1, 2, 5]
+        arg_dict["device"] = ["cpu", "cuda"]
+        for arg in GenArgList(arg_dict):
+            _test_in_top_k_impl(test_case, *arg)
+
+
+if __name__ == "__main__":
+    unittest.main()