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 StackNet(nn.Cell): 29 def __init__(self, nptype): 30 super(StackNet, self).__init__() 31 32 self.stack = P.Stack(axis=2) 33 self.data_np = np.array([0] * 16).astype(nptype) 34 self.data_np = np.reshape(self.data_np, (2, 2, 2, 2)) 35 self.x1 = Parameter(initializer( 36 Tensor(self.data_np), [2, 2, 2, 2]), name='x1') 37 self.x2 = Parameter(initializer( 38 Tensor(np.arange(16).reshape(2, 2, 2, 2).astype(nptype)), [2, 2, 2, 2]), name='x2') 39 40 @ms_function 41 def construct(self): 42 return self.stack((self.x1, self.x2)) 43 44 45def stack(nptype): 46 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 47 stack_ = StackNet(nptype) 48 output = stack_() 49 expect = np.array([[[[[0, 0], 50 [0, 0]], 51 [[0, 1], 52 [2, 3]]], 53 [[[0, 0], 54 [0, 0]], 55 [[4, 5], 56 [6, 7]]]], 57 [[[[0, 0], 58 [0, 0]], 59 [[8, 9], 60 [10, 11]]], 61 [[[0, 0], 62 [0, 0]], 63 [[12, 13], 64 [14, 15]]]]]).astype(nptype) 65 assert (output.asnumpy() == expect).all() 66 67def stack_pynative(nptype): 68 context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU') 69 x1 = np.array([0] * 16).astype(nptype) 70 x1 = np.reshape(x1, (2, 2, 2, 2)) 71 x1 = Tensor(x1) 72 x2 = Tensor(np.arange(16).reshape(2, 2, 2, 2).astype(nptype)) 73 expect = np.array([[[[[0, 0], 74 [0, 0]], 75 [[0, 1], 76 [2, 3]]], 77 [[[0, 0], 78 [0, 0]], 79 [[4, 5], 80 [6, 7]]]], 81 [[[[0, 0], 82 [0, 0]], 83 [[8, 9], 84 [10, 11]]], 85 [[[0, 0], 86 [0, 0]], 87 [[12, 13], 88 [14, 15]]]]]).astype(nptype) 89 output = P.Stack(axis=2)((x1, x2)) 90 assert (output.asnumpy() == expect).all() 91 92@pytest.mark.level0 93@pytest.mark.platform_x86_gpu_training 94@pytest.mark.env_onecard 95def test_stack_graph_float32(): 96 stack(np.float32) 97 98@pytest.mark.level0 99@pytest.mark.platform_x86_gpu_training 100@pytest.mark.env_onecard 101def test_stack_graph_float16(): 102 stack(np.float16) 103 104@pytest.mark.level1 105@pytest.mark.platform_x86_gpu_training 106@pytest.mark.env_onecard 107def test_stack_graph_int32(): 108 stack(np.int32) 109 110@pytest.mark.level1 111@pytest.mark.platform_x86_gpu_training 112@pytest.mark.env_onecard 113def test_stack_graph_int16(): 114 stack(np.int16) 115 116@pytest.mark.level1 117@pytest.mark.platform_x86_gpu_training 118@pytest.mark.env_onecard 119def test_stack_graph_uint8(): 120 stack(np.uint8) 121 122@pytest.mark.level1 123@pytest.mark.platform_x86_gpu_training 124@pytest.mark.env_onecard 125def test_stack_graph_bool(): 126 stack(np.bool) 127 128@pytest.mark.level0 129@pytest.mark.platform_x86_gpu_training 130@pytest.mark.env_onecard 131def test_stack_pynative_float32(): 132 stack_pynative(np.float32) 133 134@pytest.mark.level0 135@pytest.mark.platform_x86_gpu_training 136@pytest.mark.env_onecard 137def test_stack_pynative_float16(): 138 stack_pynative(np.float16) 139 140@pytest.mark.level1 141@pytest.mark.platform_x86_gpu_training 142@pytest.mark.env_onecard 143def test_stack_pynative_int32(): 144 stack_pynative(np.int32) 145 146@pytest.mark.level1 147@pytest.mark.platform_x86_gpu_training 148@pytest.mark.env_onecard 149def test_stack_pynative_int16(): 150 stack_pynative(np.int16) 151 152@pytest.mark.level1 153@pytest.mark.platform_x86_gpu_training 154@pytest.mark.env_onecard 155def test_stack_pynative_uint8(): 156 stack_pynative(np.uint8) 157 158@pytest.mark.level1 159@pytest.mark.platform_x86_gpu_training 160@pytest.mark.env_onecard 161def test_stack_pynative_bool(): 162 stack_pynative(np.bool) 163