use GE backend for graph mode

mindspore-lab · Jun 17, 2024 · 0396912 · 0396912
1 parent 8f060e9
commit 0396912
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 0 deletions.
diff --git a/demo/predict.py b/demo/predict.py
@@ -53,6 +53,8 @@ def get_parser_infer(parents=None):
 def set_default_infer(args):
     # Set Context
     context.set_context(mode=args.ms_mode, device_target=args.device_target, max_call_depth=2000)
+    if args.mode == 0:
+        ms.set_context(jit_config={"jit_level": "O2"})
     if args.device_target == "Ascend":
         context.set_context(device_id=int(os.getenv("DEVICE_ID", 0)))
     elif args.device_target == "GPU" and args.ms_enable_graph_kernel:

diff --git a/mindyolo/utils/utils.py b/mindyolo/utils/utils.py
@@ -22,6 +22,8 @@ def set_seed(seed=2):
 def set_default(args):
     # Set Context
     context.set_context(mode=args.ms_mode, device_target=args.device_target, max_call_depth=2000)
+    if args.mode == 0:
+        ms.set_context(jit_config={"jit_level": "O2"})
     if args.device_target == "Ascend":
         device_id = int(os.getenv("DEVICE_ID", 0))
         context.set_context(device_id=device_id)

diff --git a/test.py b/test.py
@@ -71,6 +71,8 @@ def get_parser_test(parents=None):
 def set_default_test(args):
     # Set Context
     context.set_context(mode=args.ms_mode, device_target=args.device_target, max_call_depth=2000)
+    if args.mode == 0:
+        ms.set_context(jit_config={"jit_level": "O2"})
     if args.device_target == "Ascend":
         context.set_context(device_id=int(os.getenv("DEVICE_ID", 0)))
     elif args.device_target == "GPU" and args.ms_enable_graph_kernel: