def train():
log.info('<STR_LIT>')<EOL>if version_2:<EOL><INDENT>train_data = SQuAD('<STR_LIT:train>', version='<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>train_data = SQuAD('<STR_LIT:train>', version='<STR_LIT>')<EOL><DEDENT>log.info('<STR_LIT>'.format(len(train_data)))<EOL>train_data_transform, _ = preprocess_dataset(<EOL>train_data, SQuADTransform(<EOL>berttoken,<EOL>max_seq_length=max_seq_length,<EOL>doc_stride=doc_stride,<EOL>max_query_length=max_query_length,<EOL>is_pad=True,<EOL>is_training=True))<EOL>log.info('<STR_LIT>'.format(<EOL>len(train_data_transform)))<EOL>train_dataloader = mx.gluon.data.DataLoader(<EOL>train_data_transform, batchify_fn=batchify_fn,<EOL>batch_size=batch_size, num_workers=<NUM_LIT:4>, shuffle=True)<EOL>log.info('<STR_LIT>')<EOL>optimizer_params = {'<STR_LIT>': lr}<EOL>try:<EOL><INDENT>trainer = gluon.Trainer(net.collect_params(), optimizer,<EOL>optimizer_params, update_on_kvstore=False)<EOL><DEDENT>except ValueError as e:<EOL><INDENT>print(e)<EOL>warnings.warn('<STR_LIT>'<EOL>'<STR_LIT>')<EOL>trainer = gluon.Trainer(net.collect_params(), '<STR_LIT>',<EOL>optimizer_params, update_on_kvstore=False)<EOL><DEDENT>num_train_examples = len(train_data_transform)<EOL>step_size = batch_size * accumulate if accumulate else batch_size<EOL>num_train_steps = int(num_train_examples / step_size * epochs)<EOL>num_warmup_steps = int(num_train_steps * warmup_ratio)<EOL>step_num = <NUM_LIT:0><EOL>def set_new_lr(step_num, batch_id):<EOL><INDENT>"""<STR_LIT>"""<EOL>if accumulate:<EOL><INDENT>if batch_id % accumulate == <NUM_LIT:0>:<EOL><INDENT>net.collect_params().zero_grad()<EOL>step_num += <NUM_LIT:1><EOL><DEDENT><DEDENT>else:<EOL><INDENT>step_num += <NUM_LIT:1><EOL><DEDENT>if step_num < num_warmup_steps:<EOL><INDENT>new_lr = lr * step_num / num_warmup_steps<EOL><DEDENT>else:<EOL><INDENT>offset = (step_num - num_warmup_steps) * lr /(num_train_steps - num_warmup_steps)<EOL>new_lr = lr - offset<EOL><DEDENT>trainer.set_learning_rate(new_lr)<EOL>return step_num<EOL><DEDENT>for _, v in net.collect_params('<STR_LIT>').items():<EOL><INDENT>v.wd_mult = <NUM_LIT:0.0><EOL><DEDENT>params = [p for p in net.collect_params().values()<EOL>if p.grad_req != '<STR_LIT:null>']<EOL>if accumulate:<EOL><INDENT>for p in params:<EOL><INDENT>p.grad_req = '<STR_LIT>'<EOL><DEDENT><DEDENT>epoch_tic = time.time()<EOL>total_num = <NUM_LIT:0><EOL>log_num = <NUM_LIT:0><EOL>for epoch_id in range(epochs):<EOL><INDENT>step_loss = <NUM_LIT:0.0><EOL>tic = time.time()<EOL>for batch_id, data in enumerate(train_dataloader):<EOL><INDENT>step_num = set_new_lr(step_num, batch_id)<EOL>with mx.autograd.record():<EOL><INDENT>_, inputs, token_types, valid_length, start_label, end_label = data<EOL>log_num += len(inputs)<EOL>total_num += len(inputs)<EOL>out = net(inputs.astype('<STR_LIT>').as_in_context(ctx),<EOL>token_types.astype('<STR_LIT>').as_in_context(ctx),<EOL>valid_length.astype('<STR_LIT>').as_in_context(ctx))<EOL>ls = loss_function(out, [<EOL>start_label.astype('<STR_LIT>').as_in_context(ctx),<EOL>end_label.astype('<STR_LIT>').as_in_context(ctx)]).mean()<EOL>if accumulate:<EOL><INDENT>ls = ls / accumulate<EOL><DEDENT><DEDENT>ls.backward()<EOL>if not accumulate or (batch_id + <NUM_LIT:1>) % accumulate == <NUM_LIT:0>:<EOL><INDENT>trainer.allreduce_grads()<EOL>nlp.utils.clip_grad_global_norm(params, <NUM_LIT:1>)<EOL>trainer.update(<NUM_LIT:1>)<EOL><DEDENT>step_loss += ls.asscalar()<EOL>if (batch_id + <NUM_LIT:1>) % log_interval == <NUM_LIT:0>:<EOL><INDENT>toc = time.time()<EOL>log.info(<EOL>'<STR_LIT>'  <EOL>.format(epoch_id, batch_id, len(train_dataloader),<EOL>step_loss / log_interval,<EOL>trainer.learning_rate, toc - tic, log_num / (toc - tic)))<EOL>tic = time.time()<EOL>step_loss = <NUM_LIT:0.0><EOL>log_num = <NUM_LIT:0><EOL><DEDENT><DEDENT>epoch_toc = time.time()<EOL>log.info('<STR_LIT>'<EOL>.format(epoch_id, epoch_toc - epoch_tic,<EOL>len(train_dataloader) / (epoch_toc - epoch_tic)))<EOL><DEDENT>net.save_parameters(os.path.join(output_dir, '<STR_LIT>'))<EOL>
Training function.