Browse Source

!7102 delete device id for gpu

Merge pull request !7102 from yoonlee666/tinybertbugfix
tags/v1.1.0
mindspore-ci-bot Gitee 5 years ago
parent
commit
2fcb21ec07
2 changed files with 15 additions and 2 deletions
  1. +7
    -1
      model_zoo/official/nlp/tinybert/run_general_distill.py
  2. +8
    -1
      model_zoo/official/nlp/tinybert/run_task_distill.py

+ 7
- 1
model_zoo/official/nlp/tinybert/run_general_distill.py View File

@@ -62,7 +62,13 @@ def run_general_distill():
help="dataset type tfrecord/mindrecord, default is tfrecord")
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 args_opt.device_target == "Ascend":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
elif args_opt.device_target == "GPU":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target)
else:
raise Exception("Target error, GPU or Ascend is supported.")

context.set_context(reserve_class_name_in_scope=False)
context.set_context(variable_memory_max_size="30GB")



+ 8
- 1
model_zoo/official/nlp/tinybert/run_task_distill.py View File

@@ -184,7 +184,14 @@ def run_task_distill(ckpt_file):
if ckpt_file == '':
raise ValueError("Student ckpt file should not be None")
cfg = phase2_cfg
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)

if args_opt.device_target == "Ascend":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target, device_id=args_opt.device_id)
elif args_opt.device_target == "GPU":
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.device_target)
else:
raise Exception("Target error, GPU or Ascend is supported.")

load_teacher_checkpoint_path = args_opt.load_teacher_ckpt_path
load_student_checkpoint_path = ckpt_file
netwithloss = BertNetworkWithLoss_td(teacher_config=td_teacher_net_cfg, teacher_ckpt=load_teacher_checkpoint_path,


Loading…
Cancel
Save