p60975123
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!18045 [GraphKernel]Enable Graph Kernel for FasterRCNN on GPU Device 

Merge pull request !18045 from jiaoy1224/enable
This commit is contained in:
i-robot 2021-06-10 09:51:24 +08:00 committed by Gitee
parent 0a60695fc9 f2ca5ca8d2
commit 984e353afa
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model_zoo/official/cv/faster_rcnn/train.py
View File

@ -53,6 +53,8 @@ args_opt = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
if __name__ == '__main__':
if args_opt.device_target == "GPU":
context.set_context(enable_graph_kernel=True)
if args_opt.run_distribute:
if args_opt.device_target == "Ascend":
rank = args_opt.rank_id