1# Copyright 2020 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 21import mindspore.ops.operations.array_ops as P 22from mindspore import Tensor 23from mindspore.common.api import ms_function 24from mindspore.common.initializer import initializer 25from mindspore.common.parameter import Parameter 26 27 28class UnstackNet(nn.Cell): 29 def __init__(self, nptype): 30 super(UnstackNet, self).__init__() 31 32 self.unstack = P.Unstack(axis=3) 33 self.data_np = np.array([[[[[0, 0], 34 [0, 1]], 35 [[0, 0], 36 [2, 3]]], 37 [[[0, 0], 38 [4, 5]], 39 [[0, 0], 40 [6, 7]]]], 41 [[[[0, 0], 42 [8, 9]], 43 [[0, 0], 44 [10, 11]]], 45 [[[0, 0], 46 [12, 13]], 47 [[0, 0], 48 [14, 15]]]]]).astype(nptype) 49 self.x1 = Parameter(initializer(Tensor(self.data_np), [2, 2, 2, 2, 2]), name='x1') 50 51 @ms_function 52 def construct(self): 53 return self.unstack(self.x1) 54 55 56def unstack(nptype): 57 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 58 unstack_ = UnstackNet(nptype) 59 output = unstack_() 60 expect = (np.reshape(np.array([0] * 16).astype(nptype), (2, 2, 2, 2)), 61 np.arange(2 * 2 * 2 * 2).reshape(2, 2, 2, 2).astype(nptype)) 62 63 for i, exp in enumerate(expect): 64 assert (output[i].asnumpy() == exp).all() 65 66 67def unstack_pynative(nptype): 68 context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU') 69 x1 = np.array([[[[[0, 0], 70 [0, 1]], 71 [[0, 0], 72 [2, 3]]], 73 [[[0, 0], 74 [4, 5]], 75 [[0, 0], 76 [6, 7]]]], 77 [[[[0, 0], 78 [8, 9]], 79 [[0, 0], 80 [10, 11]]], 81 [[[0, 0], 82 [12, 13]], 83 [[0, 0], 84 [14, 15]]]]]).astype(nptype) 85 x1 = Tensor(x1) 86 expect = (np.reshape(np.array([0] * 16).astype(nptype), (2, 2, 2, 2)), 87 np.arange(2 * 2 * 2 * 2).reshape(2, 2, 2, 2).astype(nptype)) 88 output = P.Unstack(axis=3)(x1) 89 90 for i, exp in enumerate(expect): 91 assert (output[i].asnumpy() == exp).all() 92 93 94@pytest.mark.level0 95@pytest.mark.platform_x86_gpu_training 96@pytest.mark.env_onecard 97def test_unstack_graph_float32(): 98 unstack(np.float32) 99 100 101@pytest.mark.level0 102@pytest.mark.platform_x86_gpu_training 103@pytest.mark.env_onecard 104def test_unstack_graph_float16(): 105 unstack(np.float16) 106 107 108@pytest.mark.level1 109@pytest.mark.platform_x86_gpu_training 110@pytest.mark.env_onecard 111def test_unstack_graph_int32(): 112 unstack(np.int32) 113 114 115@pytest.mark.level1 116@pytest.mark.platform_x86_gpu_training 117@pytest.mark.env_onecard 118def test_unstack_graph_int16(): 119 unstack(np.int16) 120 121 122@pytest.mark.level1 123@pytest.mark.platform_x86_gpu_training 124@pytest.mark.env_onecard 125def test_unstack_graph_uint8(): 126 unstack(np.uint8) 127 128 129@pytest.mark.level1 130@pytest.mark.platform_x86_gpu_training 131@pytest.mark.env_onecard 132def test_unstack_graph_bool(): 133 unstack(np.bool) 134 135 136@pytest.mark.level0 137@pytest.mark.platform_x86_gpu_training 138@pytest.mark.env_onecard 139def test_unstack_pynative_float32(): 140 unstack_pynative(np.float32) 141 142 143@pytest.mark.level0 144@pytest.mark.platform_x86_gpu_training 145@pytest.mark.env_onecard 146def test_unstack_pynative_float16(): 147 unstack_pynative(np.float16) 148 149 150@pytest.mark.level1 151@pytest.mark.platform_x86_gpu_training 152@pytest.mark.env_onecard 153def test_unstack_pynative_int32(): 154 unstack_pynative(np.int32) 155 156 157@pytest.mark.level1 158@pytest.mark.platform_x86_gpu_training 159@pytest.mark.env_onecard 160def test_unstack_pynative_int16(): 161 unstack_pynative(np.int16) 162 163 164@pytest.mark.level1 165@pytest.mark.platform_x86_gpu_training 166@pytest.mark.env_onecard 167def test_unstack_pynative_uint8(): 168 unstack_pynative(np.uint8) 169 170 171@pytest.mark.level1 172@pytest.mark.platform_x86_gpu_training 173@pytest.mark.env_onecard 174def test_unstack_pynative_bool(): 175 unstack_pynative(np.bool) 176