update jit level

examples/conformer/predict.py

@@ -5,12 +5,13 @@
 
 import os
 
+import mindspore
 import numpy as np
 from asr_model import creadte_asr_model
 from dataset import create_asr_predict_dataset, load_language_dict
-from mindspore import context
 from mindspore.train.model import Model
 from mindspore.train.serialization import load_checkpoint, load_param_into_net
+from mindspore._c_expression import ms_ctx_param
 
 from mindaudio.metric.wer import wer
 from mindaudio.models.decoders.decoder_factory import (
@@ -44,9 +45,13 @@ def main():
     os.makedirs(decode_dir, exist_ok=True)
     result_file = open(os.path.join(decode_dir, "result.txt"), "w")
 
-    context.set_context(
-        mode=context.GRAPH_MODE, device_target="Ascend", device_id=get_device_id()
+    mindspore.set_context(
+        mode=0,
+        device_target="Ascend",
+        device_id=get_device_id(),
     )
+    if "jit_config" in ms_ctx_param.__members__ and config.mode == 0:
+        mindspore.set_context(jit_config={"jit_level": "O2"})
 
     # load test data
     test_dataset = create_asr_predict_dataset(

examples/conformer/train.py

@@ -5,16 +5,18 @@
 
 import os
 
+import mindspore
 from asr_model import creadte_asr_model, create_asr_eval_net
 from dataset import create_dataset
-from mindspore import ParameterTuple, context, set_seed
+from mindspore import ParameterTuple, set_seed
 from mindspore.communication.management import init
 from mindspore.context import ParallelMode
 from mindspore.nn.optim import Adam
 from mindspore.nn.wrap.loss_scale import DynamicLossScaleUpdateCell
 from mindspore.train import Model
 from mindspore.train.callback import CheckpointConfig, ModelCheckpoint, SummaryCollector
 from mindspore.train.serialization import load_checkpoint, load_param_into_net
+from mindspore._c_expression import ms_ctx_param
 
 from mindaudio.scheduler.scheduler_factory import ASRWarmupLR
 from mindaudio.utils.callback import (
@@ -57,21 +59,23 @@ def train():
     model_dir = os.path.join(exp_dir, "model")
     graph_dir = os.path.join(exp_dir, "graph")
     summary_dir = os.path.join(exp_dir, "summary")
-    context.set_context(
-        mode=context.GRAPH_MODE,
+    mindspore.set_context(
+        mode=0,
         device_target="Ascend",
         device_id=get_device_id(),
         save_graphs=config.save_graphs,
         save_graphs_path=graph_dir,
     )
+    if "jit_config" in ms_ctx_param.__members__ and config.mode == 0:
+        mindspore.set_context(jit_config={"jit_level": "O2"})
 
     device_num = get_device_num()
     rank = get_rank_id()
     # configurations for distributed training
     if config.is_distributed:
         init()
-        context.reset_auto_parallel_context()
-        context.set_auto_parallel_context(
+        mindspore.reset_auto_parallel_context()
+        mindspore.set_auto_parallel_context(
             parallel_mode=ParallelMode.DATA_PARALLEL,
             gradients_mean=True,
             device_num=device_num,

requirements-dev.txt

@@ -1,5 +1,5 @@
 mindspore==2.0.0
-numpy>=1.17.0
+numpy>=1.17.0, <2
 scipy>=1.6.0
 pyyaml>=5.3
 tqdm

requirements.txt

@@ -1,5 +1,5 @@
 mindspore==2.0.0
-numpy>=1.17.0
+numpy>=1.17.0, <2
 scipy>=1.6.0
 pyyaml>=5.3
 tqdm