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""" test loss """ 16import numpy as np 17import pytest 18 19import mindspore as ms 20import mindspore.nn as nn 21import mindspore.context as context 22from mindspore import Tensor 23from mindspore.ops import operations as P 24 25def func_single_output(x1, x2): 26 return x1 - x2 27 28def func_multi_output(x1, x2): 29 return (x1 + x2), (x1 - x2) 30 31output = 0 32def func_no_output(x1, x2): 33 global output 34 output = x1 + x2 35 36class PyFuncNet(nn.Cell): 37 def __init__(self, fn, in_types, in_shapes, out_types, out_shapes): 38 super().__init__() 39 self.func = P.PyFunc(fn, in_types, in_shapes, out_types, out_shapes) 40 self.relu = P.ReLU() 41 42 def construct(self, x1, x2): 43 x = self.func((x1, x2)) 44 return self.relu(x[0]) 45 46 47def func_with_dtype(ms_dtype, np_dtype): 48 shape = (40, 40) 49 np.random.seed(42) 50 x1 = np.random.randint(-5, 5, size=shape).astype(np_dtype) 51 x2 = np.random.randint(-5, 5, size=shape).astype(np_dtype) 52 53 expect = func_single_output(x1, x2) 54 expect = P.ReLU()(Tensor(expect)) 55 56 net = PyFuncNet(func_single_output, [ms_dtype, ms_dtype], [shape, shape], [ms_dtype], [shape]) 57 x = net(Tensor(x1), Tensor(x2)) 58 assert np.allclose(x.asnumpy(), expect.asnumpy()) 59 60 61@pytest.mark.level0 62@pytest.mark.platform_x86_gpu_training 63@pytest.mark.env_onecard 64def test_pyfunc_single_output(): 65 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 66 func_with_dtype(ms.float16, np.float16) 67 func_with_dtype(ms.float32, np.float32) 68 func_with_dtype(ms.float64, np.float64) 69 func_with_dtype(ms.int32, np.int32) 70 func_with_dtype(ms.int64, np.int64) 71 72 73@pytest.mark.level0 74@pytest.mark.platform_x86_gpu_training 75@pytest.mark.env_onecard 76def test_pyfunc_multi_output(): 77 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 78 shape = (40, 40) 79 dtype = ms.float32 80 81 np.random.seed(42) 82 x1 = np.random.randint(-5, 5, size=shape).astype(np.float32) 83 x2 = np.random.randint(-5, 5, size=shape).astype(np.float32) 84 expect, _ = func_multi_output(x1, x2) 85 expect = P.ReLU()(Tensor(expect)) 86 87 net = PyFuncNet(func_multi_output, [dtype, dtype], [shape, shape], [dtype, dtype], [shape, shape]) 88 x = net(Tensor(x1), Tensor(x2)) 89 90 assert np.allclose(x.asnumpy(), expect.asnumpy()) 91 92 93class PyFuncGraph(nn.Cell): 94 def __init__(self, fn, in_types, in_shapes, out_types, out_shapes): 95 super().__init__() 96 self.func = P.PyFunc(fn, in_types, in_shapes, out_types, out_shapes) 97 98 def construct(self, x1, x2): 99 return self.func((x1, x2)) 100 101@pytest.mark.level0 102@pytest.mark.platform_x86_gpu_training 103@pytest.mark.env_onecard 104def test_pyfunc_no_output(): 105 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 106 shape = (40, 40) 107 dtype = ms.float32 108 109 np.random.seed(42) 110 x1 = np.random.randint(-5, 5, size=shape).astype(np.float32) 111 x2 = np.random.randint(-5, 5, size=shape).astype(np.float32) 112 func_no_output(x1, x2) 113 global output 114 expect = output 115 116 net = PyFuncGraph(func_no_output, [dtype, dtype], [shape, shape], [], []) 117 net(Tensor(x1), Tensor(x2)) 118 net_output = output 119 120 assert np.allclose(net_output, expect) 121 122 123@pytest.mark.level0 124@pytest.mark.platform_x86_gpu_training 125@pytest.mark.env_onecard 126def test_pyfunc_scalar(): 127 context.set_context(mode=context.GRAPH_MODE, device_target='GPU') 128 shape = () 129 ms_dtype = ms.int32 130 131 x1 = int(10) 132 x2 = int(5) 133 expect = func_single_output(x1, x2) 134 135 net = PyFuncGraph(func_single_output, [ms_dtype, ms_dtype], [shape, shape], [ms_dtype], [shape]) 136 x = net(Tensor(x1), Tensor(x2)) 137 assert np.allclose(x[0].asnumpy(), expect) 138