“8a24b8993dfed89f3c7d56d67c6c132ada2e5bc9”上不存在“tests/pytest/git@gitcode.net:qq_37101384/tdengine.git”
test_top_k_op.py 3.4 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# 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.

15 16
from __future__ import print_function

武毅 已提交
17 18
import unittest
import numpy as np
19
from op_test import OpTest
武毅 已提交
20 21


Q
qijun 已提交
22
class TestTopkOp(OpTest):
武毅 已提交
23
    def setUp(self):
W
whs 已提交
24
        self.variable_k = False
Q
qingqing01 已提交
25
        self.set_args()
Q
qijun 已提交
26
        self.op_type = "top_k"
W
Wu Yi 已提交
27 28 29
        self.dtype = np.float32
        self.init_dtype()

Q
qingqing01 已提交
30
        k = self.top_k
W
Wu Yi 已提交
31
        input = np.random.random((self.row, k)).astype(self.dtype)
Q
qingqing01 已提交
32 33
        output = np.ndarray((self.row, k))
        indices = np.ndarray((self.row, k)).astype("int64")
武毅 已提交
34
        self.inputs = {'X': input}
W
whs 已提交
35 36 37 38 39

        if self.variable_k:
            self.inputs['K'] = np.array([k]).astype("int32")
        else:
            self.attrs = {'k': k}
武毅 已提交
40

Q
qingqing01 已提交
41
        for rowid in range(self.row):
武毅 已提交
42
            row = input[rowid]
Q
qingqing01 已提交
43 44
            output[rowid] = np.sort(row)[::-1][:k]
            indices[rowid] = row.argsort()[::-1][:k]
武毅 已提交
45 46 47

        self.outputs = {'Out': output, 'Indices': indices}

W
Wu Yi 已提交
48 49 50
    def init_dtype(self):
        pass

Q
qingqing01 已提交
51 52 53 54
    def set_args(self):
        self.row = 32
        self.top_k = 1

55 56 57
    def test_check_output(self):
        self.check_output()

武毅 已提交
58

W
Wu Yi 已提交
59 60 61 62 63
class TestTopkOpFp16(TestTopkOp):
    def init_dtype(self):
        self.dtype = np.float16


Q
qijun 已提交
64
class TestTopkOp3d(OpTest):
武毅 已提交
65
    def setUp(self):
Q
qijun 已提交
66
        self.op_type = "top_k"
武毅 已提交
67 68 69 70
        k = 1
        input = np.random.random((32, 2, 84)).astype("float32")
        input_flat_2d = input.reshape(64, 84)
        output = np.ndarray((64, k))
F
fengjiayi 已提交
71
        indices = np.ndarray((64, k)).astype("int64")
武毅 已提交
72

Q
qingqing01 已提交
73
        self.inputs = {'X': input}
武毅 已提交
74 75
        self.attrs = {'k': k}

76
        for rowid in range(64):
武毅 已提交
77
            row = input_flat_2d[rowid]
Q
qingqing01 已提交
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105
            output[rowid] = np.sort(row)[::-1][:k]
            indices[rowid] = row.argsort()[::-1][:k]

        self.outputs = {
            'Out': output.reshape((32, 2, k)),
            'Indices': indices.reshape((32, 2, k))
        }

    def test_check_output(self):
        self.check_output()


class TestTopkOp2(OpTest):
    def setUp(self):
        self.op_type = "top_k"
        k = 1
        m = 2056
        input = np.random.random((m, 84)).astype("float32")
        output = np.ndarray((m, k))
        indices = np.ndarray((m, k)).astype("int64")

        self.inputs = {'X': input}
        self.attrs = {'k': k}

        for rowid in range(m):
            row = input[rowid]
            output[rowid] = -np.sort(-row)[:k]
            indices[rowid] = (-row).argsort()[:k]
武毅 已提交
106 107 108

        self.outputs = {'Out': output, 'Indices': indices}

109 110 111
    def test_check_output(self):
        self.check_output()

武毅 已提交
112

Q
qingqing01 已提交
113 114 115 116 117 118 119 120 121 122 123 124
class TestTopkOp3(TestTopkOp):
    def set_args(self):
        self.row = 2056
        self.top_k = 3


class TestTopkOp4(TestTopkOp):
    def set_args(self):
        self.row = 40000
        self.top_k = 1


W
whs 已提交
125 126 127 128 129 130 131
class TestTopkOp5(TestTopkOp):
    def set_args(self):
        self.row = 40000
        self.top_k = 3
        self.variable_k = True


Q
qijun 已提交
132
if __name__ == "__main__":
武毅 已提交
133
    unittest.main()