From 43b13f4a500a7c8901601cf0f476f06580721a5f Mon Sep 17 00:00:00 2001 From: Margaret_wangrui Date: Fri, 23 Apr 2021 11:44:49 +0800 Subject: [PATCH] add control flow testcases --- tests/st/control/inner/test_002_single_for.py | 104 ++++++++++++++++++ tests/st/control/inner/test_030_for_in_if.py | 52 +++++++++ 2 files changed, 156 insertions(+) diff --git a/tests/st/control/inner/test_002_single_for.py b/tests/st/control/inner/test_002_single_for.py index b357058089..62649f17d2 100644 --- a/tests/st/control/inner/test_002_single_for.py +++ b/tests/st/control/inner/test_002_single_for.py @@ -12,8 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ +import numpy as np from mindspore import context from mindspore import Tensor, nn +from mindspore.common.parameter import Parameter from mindspore.ops import composite as C from mindspore.ops import operations as P from mindspore.common import dtype as mstype @@ -107,3 +109,105 @@ def test_single_for_02(): assert graph_forward_res == pynative_forward_res assert graph_backward_res == pynative_backward_res + + +def test_single_for_03(): + class SingleForNet(nn.Cell): + def __init__(self): + super().__init__() + self.mul = P.Mul() + self.add = P.Add() + self.sub = P.Sub() + self.assign = P.Assign() + param_a = np.full((1,), 5, dtype=np.float32) + self.param_a = Parameter(Tensor(param_a), name='a') + param_b = np.full((1,), 2, dtype=np.float32) + self.param_b = Parameter(Tensor(param_b), name='b') + + def func(self, x): + x = self.mul(x, 2) + for _ in range(0, 5): + x = self.add(x, x) + self.param_b = self.param_b + 1 + return x - self.param_b + + def construct(self, x, y): + self.assign(self.param_a, x + self.param_a) + z = self.func(x) + x = self.param_a + y + z + return x, self.param_b + + class GradNet(nn.Cell): + def __init__(self, net): + super(GradNet, self).__init__() + self.net = net + + def construct(self, *inputs): + return grad_all(self.net)(*inputs) + + x = Tensor([2], mstype.int32) + y = Tensor([5], mstype.int32) + + # graph mode + context.set_context(mode=context.GRAPH_MODE) + single_for_net = SingleForNet() + net = GradNet(single_for_net) + graph_forward_res = single_for_net(x, y) + graph_backward_res = net(x, y) + + # pynative mode + context.set_context(mode=context.PYNATIVE_MODE) + single_for_net = SingleForNet() + net = GradNet(single_for_net) + pynative_forward_res = single_for_net(x, y) + pynative_backward_res = net(x, y) + + assert graph_forward_res == pynative_forward_res + assert graph_backward_res == pynative_backward_res + + +def test_single_for_04(): + class SingleForNet(nn.Cell): + def __init__(self): + super().__init__() + self.mul = P.Mul() + self.add = P.Add() + self.sub = P.Sub() + self.assign = P.Assign() + param_a = np.full((1,), 5, dtype=np.float32) + self.param_a = Parameter(Tensor(param_a), name='a') + param_b = np.full((1,), 2, dtype=np.float32) + self.param_b = Parameter(Tensor(param_b), name='b') + + def construct(self, x): + self.assign(self.param_a, x + self.param_a) + for _ in range(1): + self.param_b = x - self.param_a + return self.param_b + + class GradNet(nn.Cell): + def __init__(self, net): + super(GradNet, self).__init__() + self.net = net + + def construct(self, *inputs): + return grad_all(self.net)(*inputs) + + x = Tensor([2], mstype.int32) + + # graph mode + context.set_context(mode=context.GRAPH_MODE) + single_for_net = SingleForNet() + net = GradNet(single_for_net) + graph_forward_res = single_for_net(x) + graph_backward_res = net(x) + + # pynative mode + context.set_context(mode=context.PYNATIVE_MODE) + single_for_net = SingleForNet() + net = GradNet(single_for_net) + pynative_forward_res = single_for_net(x) + pynative_backward_res = net(x) + + assert graph_forward_res == pynative_forward_res + assert graph_backward_res == pynative_backward_res diff --git a/tests/st/control/inner/test_030_for_in_if.py b/tests/st/control/inner/test_030_for_in_if.py index 721f200ed1..6b01497253 100644 --- a/tests/st/control/inner/test_030_for_in_if.py +++ b/tests/st/control/inner/test_030_for_in_if.py @@ -118,3 +118,55 @@ def test_for_in_if_02(): assert graph_forward_res == pynative_forward_res assert graph_backward_res == pynative_backward_res + + +def test_for_in_if_03(): + class ForInIfNet(nn.Cell): + def __init__(self): + super().__init__() + self.mul = P.Mul() + self.add = P.Add() + param_a = np.full((1,), 5, dtype=np.float32) + self.param_a = Parameter(Tensor(param_a), name='a') + param_b = np.full((1,), 4, dtype=np.float32) + self.param_b = Parameter(Tensor(param_b), name='b') + + def construct(self, x): + y = x + self.param_b + if self.param_a > self.param_b: + x = self.mul(x, 2) + for i in range(-1, 5): + x = self.add(i, x) + self.param_b += 1 + elif y > x: + y = self.param_a * y + else: + x = self.param_b * x + return x, y + + class GradNet(nn.Cell): + def __init__(self, net): + super(GradNet, self).__init__() + self.net = net + + def construct(self, *inputs): + return grad_all(self.net)(*inputs) + + x = Tensor([10], mstype.int32) + + # graph mode + context.set_context(mode=context.GRAPH_MODE) + for_in_if_net = ForInIfNet() + net = GradNet(for_in_if_net) + graph_forward_res = for_in_if_net(x) + graph_backward_res = net(x) + + # pynative mode + context.set_context(mode=context.PYNATIVE_MODE) + for_in_if_net = ForInIfNet() + net = GradNet(for_in_if_net) + pynative_forward_res = for_in_if_net(x) + pynative_backward_res = net(x) + + assert graph_forward_res == pynative_forward_res + assert graph_backward_res == pynative_backward_res