def train(metric):
logging.info('<STR_LIT>', ctx)<EOL>optimizer_params = {'<STR_LIT>': lr, '<STR_LIT>': epsilon, '<STR_LIT>': <NUM_LIT>}<EOL>try:<EOL><INDENT>trainer = gluon.Trainer(<EOL>model.collect_params(),<EOL>args.optimizer,<EOL>optimizer_params,<EOL>update_on_kvstore=False)<EOL><DEDENT>except ValueError as e:<EOL><INDENT>print(e)<EOL>warnings.warn(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL>trainer = gluon.Trainer(<EOL>model.collect_params(),<EOL>'<STR_LIT>',<EOL>optimizer_params,<EOL>update_on_kvstore=False)<EOL><DEDENT>step_size = batch_size * accumulate if accumulate else batch_size<EOL>num_train_steps = int(num_train_examples / step_size * args.epochs)<EOL>warmup_ratio = args.warmup_ratio<EOL>num_warmup_steps = int(num_train_steps * warmup_ratio)<EOL>step_num = <NUM_LIT:0><EOL>for _, v in model.collect_params('<STR_LIT>').items():<EOL><INDENT>v.wd_mult = <NUM_LIT:0.0><EOL><DEDENT>params = [<EOL>p for p in model.collect_params().values() if p.grad_req != '<STR_LIT:null>'<EOL>]<EOL>if accumulate:<EOL><INDENT>for p in params:<EOL><INDENT>p.grad_req = '<STR_LIT>'<EOL><DEDENT><DEDENT>for epoch_id in range(args.epochs):<EOL><INDENT>metric.reset()<EOL>step_loss = <NUM_LIT:0><EOL>tic = time.time()<EOL>for batch_id, seqs in enumerate(train_data):<EOL><INDENT>if accumulate:<EOL><INDENT>if batch_id % accumulate == <NUM_LIT:0>:<EOL><INDENT>model.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 / (<EOL>num_train_steps - num_warmup_steps)<EOL>new_lr = lr - offset<EOL><DEDENT>trainer.set_learning_rate(new_lr)<EOL>with mx.autograd.record():<EOL><INDENT>input_ids, valid_length, type_ids, label = seqs<EOL>out = model(<EOL>input_ids.as_in_context(ctx), type_ids.as_in_context(ctx),<EOL>valid_length.astype('<STR_LIT>').as_in_context(ctx))<EOL>ls = loss_function(out, label.as_in_context(ctx)).mean()<EOL><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(accumulate if accumulate else <NUM_LIT:1>)<EOL><DEDENT>step_loss += ls.asscalar()<EOL>metric.update([label], [out])<EOL>if (batch_id + <NUM_LIT:1>) % (args.log_interval) == <NUM_LIT:0>:<EOL><INDENT>log_train(batch_id, len(train_data), metric, step_loss, args.log_interval,<EOL>epoch_id, trainer.learning_rate)<EOL>step_loss = <NUM_LIT:0><EOL><DEDENT><DEDENT>mx.nd.waitall()<EOL>if task.task_name == '<STR_LIT>':<EOL><INDENT>logging.info('<STR_LIT>')<EOL>evaluate(dev_data_matched, metric)<EOL>logging.info('<STR_LIT>')<EOL>evaluate(dev_data_mismatched, metric)<EOL><DEDENT>else:<EOL><INDENT>evaluate(dev_data, metric)<EOL><DEDENT>params_saved = os.path.join(output_dir,<EOL>'<STR_LIT>'.format(task.task_name, epoch_id))<EOL>model.save_parameters(params_saved)<EOL>logging.info('<STR_LIT>', params_saved)<EOL>toc = time.time()<EOL>logging.info('<STR_LIT>', toc - tic)<EOL>tic = toc<EOL><DEDENT>
Training function.