• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2021 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15
16import numpy as np
17import pytest
18
19import mindspore.context as context
20import mindspore.nn as nn
21from mindspore import Tensor
22from mindspore.common.api import ms_function
23
24context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
25
26
27class NetOneHot(nn.Cell):
28    def __init__(self):
29        super(NetOneHot, self).__init__()
30        self.on_value = 2.0
31        self.off_value = 3.0
32
33        self.depth_1 = 6
34        self.one_hot_1 = nn.OneHot(-1, self.depth_1, self.on_value, self.off_value)
35
36        self.depth_2 = 4
37        self.one_hot_2 = nn.OneHot(0, self.depth_1, self.on_value, self.off_value)
38        self.one_hot_3 = nn.OneHot(0, self.depth_2, self.on_value, self.off_value)
39        self.one_hot_4 = nn.OneHot(1, self.depth_1, self.on_value, self.off_value)
40
41    @ms_function
42    def construct(self, indices1, indices2, indices3, indices4):
43        return (self.one_hot_1(indices1), self.one_hot_2(indices2),
44                self.one_hot_3(indices3), self.one_hot_4(indices4))
45
46
47def one_hot(nptype):
48    one_hot_net = NetOneHot()
49    indices1 = Tensor(np.array([[0, 1], [4, 5], [2, 6]]).astype(nptype))
50    indices2 = Tensor(np.array([1, 2, 3]).astype(nptype))
51    indices3 = Tensor(np.array([[0, 1], [1, 0]]).astype(nptype))
52    indices4 = Tensor(np.array([[0, 1], [4, 5], [2, 6]]).astype(nptype))
53    output = one_hot_net(indices1, indices2, indices3, indices4)
54    expect_0 = np.array([
55        [[2., 3., 3., 3., 3., 3.], [3., 2., 3., 3., 3., 3.]],
56        [[3., 3., 3., 3., 2., 3.], [3., 3., 3., 3., 3., 2.]],
57        [[3., 3., 2., 3., 3., 3.], [3., 3., 3., 3., 3., 3.]]
58    ]).astype(np.float32)
59    expect_1 = np.array([
60        [3., 3., 3.],
61        [2., 3., 3.],
62        [3., 2., 3.],
63        [3., 3., 2.],
64        [3., 3., 3.],
65        [3., 3., 3.]
66    ]).astype(np.float32)
67    expect_2 = np.array([
68        [[2., 3.], [3., 2.]], [[3., 2.], [2., 3.]], [[3., 3.], [3., 3.]],
69        [[3., 3.], [3., 3.]]
70    ]).astype(np.float32)
71    expect_3 = np.array([
72        [[2., 3.], [3., 2.], [3., 3.], [3., 3.], [3., 3.], [3., 3.]],
73        [[3., 3.], [3., 3.], [3., 3.], [3., 3.], [2., 3.], [3., 2.]],
74        [[3., 3.], [3., 3.], [2., 3.], [3., 3.], [3., 3.], [3., 3.]]
75    ]).astype(np.float32)
76    assert (output[0].asnumpy() == expect_0).all()
77    assert (output[1].asnumpy() == expect_1).all()
78    assert (output[2].asnumpy() == expect_2).all()
79    assert (output[3].asnumpy() == expect_3).all()
80
81@pytest.mark.level0
82@pytest.mark.platform_x86_gpu_training
83@pytest.mark.env_onecard
84def test_one_hot_int32():
85    one_hot(np.int32)
86
87@pytest.mark.level0
88@pytest.mark.platform_x86_gpu_training
89@pytest.mark.env_onecard
90def test_one_hot_int64():
91    one_hot(np.int64)
92