def train(data_train, model, nsp_loss, mlm_loss, vocab_size, ctx, store):
mlm_metric = nlp.metric.MaskedAccuracy()<EOL>nsp_metric = nlp.metric.MaskedAccuracy()<EOL>mlm_metric.reset()<EOL>nsp_metric.reset()<EOL>lr = args.lr<EOL>optim_params = {'<STR_LIT>': lr, '<STR_LIT>': <NUM_LIT>, '<STR_LIT>': <NUM_LIT>}<EOL>if args.dtype == '<STR_LIT>':<EOL><INDENT>optim_params['<STR_LIT>'] = True<EOL><DEDENT>trainer = mx.gluon.Trainer(model.collect_params(), '<STR_LIT>', optim_params,<EOL>update_on_kvstore=False, kvstore=store)<EOL>dynamic_loss_scale = args.dtype == '<STR_LIT>'<EOL>fp16_trainer = FP16Trainer(trainer, dynamic_loss_scale=dynamic_loss_scale)<EOL>if args.ckpt_dir and args.start_step:<EOL><INDENT>state_path = os.path.join(args.ckpt_dir, '<STR_LIT>' % args.start_step)<EOL>logging.info('<STR_LIT>', state_path)<EOL>trainer.load_states(state_path)<EOL><DEDENT>accumulate = args.accumulate<EOL>num_train_steps = args.num_steps<EOL>warmup_ratio = args.warmup_ratio<EOL>num_warmup_steps = int(num_train_steps * warmup_ratio)<EOL>params = [p for p in model.collect_params().values() if p.grad_req != '<STR_LIT:null>']<EOL>param_dict = model.collect_params()<EOL>for _, v in model.collect_params('<STR_LIT>').items():<EOL><INDENT>v.wd_mult = <NUM_LIT:0.0><EOL><DEDENT>if accumulate > <NUM_LIT:1>:<EOL><INDENT>for p in params:<EOL><INDENT>p.grad_req = '<STR_LIT>'<EOL><DEDENT><DEDENT>train_begin_time = time.time()<EOL>begin_time = time.time()<EOL>running_mlm_loss = running_nsp_loss = running_num_tks = <NUM_LIT:0><EOL>batch_num = <NUM_LIT:0><EOL>step_num = args.start_step<EOL>parallel_model = ParallelBERT(model, mlm_loss, nsp_loss, vocab_size,<EOL>store.num_workers * accumulate, trainer=fp16_trainer)<EOL>num_ctxes = len(ctx)<EOL>parallel = nlp.utils.Parallel(num_ctxes if num_ctxes > <NUM_LIT:1> else <NUM_LIT:0>, parallel_model)<EOL>while step_num < num_train_steps:<EOL><INDENT>for _, dataloader in enumerate(data_train):<EOL><INDENT>if step_num >= num_train_steps:<EOL><INDENT>break<EOL><DEDENT>if args.dummy_data_len:<EOL><INDENT>target_shape = (args.batch_size*num_ctxes, args.dummy_data_len)<EOL>dataloader = get_dummy_dataloader(dataloader, target_shape)<EOL><DEDENT>for _, data_batch in enumerate(dataloader):<EOL><INDENT>if step_num >= num_train_steps:<EOL><INDENT>break<EOL><DEDENT>if batch_num % accumulate == <NUM_LIT:0>:<EOL><INDENT>step_num += <NUM_LIT:1><EOL>if accumulate > <NUM_LIT:1>:<EOL><INDENT>param_dict.zero_grad()<EOL><DEDENT>if step_num <= num_warmup_steps:<EOL><INDENT>new_lr = lr * step_num / num_warmup_steps<EOL><DEDENT>else:<EOL><INDENT>offset = lr * step_num / num_train_steps<EOL>new_lr = lr - offset<EOL><DEDENT>trainer.set_learning_rate(new_lr)<EOL>if args.profile:<EOL><INDENT>profile(step_num, <NUM_LIT:10>, <NUM_LIT:12>, profile_name=args.profile)<EOL><DEDENT><DEDENT>if args.use_avg_len:<EOL><INDENT>data_list = [[seq.as_in_context(context) for seq in shard]<EOL>for context, shard in zip(ctx, data_batch)]<EOL><DEDENT>else:<EOL><INDENT>if data_batch[<NUM_LIT:0>].shape[<NUM_LIT:0>] < len(ctx):<EOL><INDENT>continue<EOL><DEDENT>data_list = split_and_load(data_batch, ctx)<EOL><DEDENT>ns_label_list, ns_pred_list = [], []<EOL>mask_label_list, mask_pred_list, mask_weight_list = [], [], []<EOL>for data in data_list:<EOL><INDENT>parallel.put(data)<EOL><DEDENT>for _ in range(len(ctx)):<EOL><INDENT>(_, next_sentence_label, classified, masked_id,<EOL>decoded, masked_weight, ls1, ls2, valid_length) = parallel.get()<EOL>ns_label_list.append(next_sentence_label)<EOL>ns_pred_list.append(classified)<EOL>mask_label_list.append(masked_id)<EOL>mask_pred_list.append(decoded)<EOL>mask_weight_list.append(masked_weight)<EOL>running_mlm_loss += ls1.as_in_context(mx.cpu()) / num_ctxes<EOL>running_nsp_loss += ls2.as_in_context(mx.cpu()) / num_ctxes<EOL>running_num_tks += valid_length.sum().as_in_context(mx.cpu())<EOL><DEDENT>if (batch_num + <NUM_LIT:1>) % accumulate == <NUM_LIT:0>:<EOL><INDENT>fp16_trainer.step(<NUM_LIT:1>, max_norm=<NUM_LIT:1>)<EOL><DEDENT>nsp_metric.update(ns_label_list, ns_pred_list)<EOL>mlm_metric.update(mask_label_list, mask_pred_list, mask_weight_list)<EOL>if (step_num + <NUM_LIT:1>) % (args.log_interval) == <NUM_LIT:0> and (batch_num + <NUM_LIT:1>) % accumulate == <NUM_LIT:0>:<EOL><INDENT>log(begin_time, running_num_tks, running_mlm_loss / accumulate,<EOL>running_nsp_loss / accumulate, step_num, mlm_metric,<EOL>nsp_metric, trainer, args.log_interval)<EOL>begin_time = time.time()<EOL>running_mlm_loss = running_nsp_loss = running_num_tks = <NUM_LIT:0><EOL>mlm_metric.reset_local()<EOL>nsp_metric.reset_local()<EOL><DEDENT>if args.ckpt_dir and (step_num + <NUM_LIT:1>) % (args.ckpt_interval) == <NUM_LIT:0>and (batch_num + <NUM_LIT:1>) % accumulate == <NUM_LIT:0>:<EOL><INDENT>save_params(step_num, model, trainer, args.ckpt_dir)<EOL><DEDENT>batch_num += <NUM_LIT:1><EOL><DEDENT><DEDENT><DEDENT>save_params(step_num, model, trainer, args.ckpt_dir)<EOL>mx.nd.waitall()<EOL>train_end_time = time.time()<EOL>logging.info('<STR_LIT>'.format(train_end_time - train_begin_time))<EOL>
Training function.