From b3a558d7f54e987aad9348c55995c53baf86a014 Mon Sep 17 00:00:00 2001 From: brian Date: Fri, 7 Feb 2025 16:36:58 +0800 Subject: [PATCH] [mod]: mod jitlevel --- tests/st/mindflow/networks/burgers/test_burgers.py | 4 ++-- .../networks/navier_stokes/test_mindflow_navier_stokes.py | 8 +++++--- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/st/mindflow/networks/burgers/test_burgers.py b/tests/st/mindflow/networks/burgers/test_burgers.py index 1d60c0b40..14776d3a7 100644 --- a/tests/st/mindflow/networks/burgers/test_burgers.py +++ b/tests/st/mindflow/networks/burgers/test_burgers.py @@ -80,7 +80,7 @@ class Net(nn.Cell): return self.layers(x) -@pytest.mark.level1 +@pytest.mark.level0 @pytest.mark.platform_arm_ascend910b_training @pytest.mark.env_onecard def test_mindflow_burgers_pinns(): @@ -89,7 +89,7 @@ def test_mindflow_burgers_pinns(): Description: test train and eval Expectation: success """ - context.set_context(mode=context.GRAPH_MODE) + context.set_context(mode=context.GRAPH_MODE, jit_config={"jit_level": "O2"}) model = Net() optimizer = nn.Adam(model.trainable_params(), 0.0001) problem = Burgers1D(model) diff --git a/tests/st/mindflow/networks/navier_stokes/test_mindflow_navier_stokes.py b/tests/st/mindflow/networks/navier_stokes/test_mindflow_navier_stokes.py index 128e75c3b..0e6e71e02 100644 --- a/tests/st/mindflow/networks/navier_stokes/test_mindflow_navier_stokes.py +++ b/tests/st/mindflow/networks/navier_stokes/test_mindflow_navier_stokes.py @@ -29,6 +29,7 @@ np.random.seed(123456) class Net(nn.Cell): """MLP""" + def __init__(self, in_channels=2, hidden_channels=128, out_channels=1, act=nn.Tanh()): super().__init__() self.act = act @@ -43,7 +44,7 @@ class Net(nn.Cell): return self.layers(x) -@pytest.mark.level1 +@pytest.mark.level0 @pytest.mark.platform_arm_ascend910b_training @pytest.mark.env_onecard def test_mindflow_navier_stokes(): @@ -52,7 +53,7 @@ def test_mindflow_navier_stokes(): Description: test train Expectation: success """ - context.set_context(mode=context.GRAPH_MODE) + context.set_context(mode=context.GRAPH_MODE, jit_config={"jit_level": "O2"}) model = Net(in_channels=3, out_channels=3) optimizer = nn.Adam(model.trainable_params(), 0.0001) problem = NavierStokes2D(model) @@ -115,7 +116,8 @@ def test_mindflow_navier_stokes(): time_beg = time.time() train_loss = train_step(pde_data, bc_data, bc_label, ic_data, ic_label) epoch_time = time.time() - time_beg - print(f"epoch: {epoch} train loss: {train_loss} epoch time: {epoch_time}s") + print( + f"epoch: {epoch} train loss: {train_loss} epoch time: {epoch_time}s") model.set_train(False) assert epoch_time < 0.01 -- Gitee